@@ -11,7 +11,7 @@ import Control.Monad.Random
1111import Control.Monad.Trans.Except
1212
1313import Data.Char ( isUpper, toUpper, toLower )
14−import Data.List ( unfoldr, foldl' )
14+import Data.List ( foldl' )
1515import Data.Maybe ( fromMaybe )
1616
1717import qualified Data.Vector as V
@@ -20,6 +20,8 @@ import Data.Vector ( Vector )
2020import qualified Data.Map as M
2121import Data.Proxy ( Proxy (..) )
2222
23+import qualified Data.ByteString as B
24+import Data.Serialize
2325
2426import Data.Singletons.Prelude
2527import GHC.TypeLits
@@ -32,29 +34,31 @@ import Grenade
3234import Grenade.Recurrent
3335import Grenade.Utils.OneHot
3436
37+import System.IO.Unsafe ( unsafeInterleaveIO )
38+
3539-- The defininition for our natural language recurrent network.
3640-- This network is able to learn and generate simple words in
3741-- about an hour.
3842--
39−-- This is a first class recurrent net, although it's similar to
40−-- an unrolled graph.
43+-- This is a first class recurrent net.
4144--
4245-- The F and R types are tagging types to ensure that the runner and
4346-- creation function know how to treat the layers.
4447--
4548-- As an example, here's a short sequence generated.
4649--
47−-- > the see and and the sir, and and the make and the make and go the make and go the make and the
48−--
50+-- > KING RICHARD III:
51+-- > And as the heaven her his words, we the son, I show sand stape but the lament to shall were the sons with a strend
52+
4953type F = FeedForward
5054type R = Recurrent
5155
5256-- The definition of our network
53−type Shakespeare = RecurrentNetwork '[ R (LSTM 40 50), R (LSTM 50 40), F (FullyConnected 40 40), F Logit]
54− '[ 'D1 40, 'D1 50, 'D1 40, 'D1 40, 'D1 40 ]
57+type Shakespeare = RecurrentNetwork '[ R (LSTM 40 80), R (LSTM 80 40), F (FullyConnected 40 40), F Logit]
58+ '[ 'D1 40, 'D1 80, 'D1 40, 'D1 40, 'D1 40 ]
5559
5660-- The definition of the "sideways" input, which the network if fed recurrently.
57−type Shakespearian = RecurrentInputs '[ R (LSTM 40 50), R (LSTM 50 40), F (FullyConnected 40 40), F Logit]
61+type Shakespearian = RecurrentInputs '[ R (LSTM 40 80), R (LSTM 80 40), F (FullyConnected 40 40), F Logit]
5862
5963randomNet :: MonadRandom m => m (Shakespeare, Shakespearian)
6064randomNet = randomRecurrent
@@ -82,15 +86,23 @@ trainSlice !rate !net !recIns input offset size =
8286runShakespeare :: ShakespeareOpts -> ExceptT String IO ()
8387runShakespeare ShakespeareOpts {..} = do
8488 (shakespeare, oneHotMap, oneHotDictionary) <- loadShakespeare trainingFile
85− (net0, i0) <- lift randomNet
86− lift $ foldM_ (\(!net, !io) size -> do
89+ (net0, i0) <- lift $
90+ case loadPath of
91+ Just loadFile -> netLoad loadFile
92+ Nothing -> randomNet
93+
94+ (trained, bestInput) <- lift $ foldM (\(!net, !io) size -> do
8795 xs <- take (iterations `div` 15) <$> getRandomRs (0, length shakespeare - size - 1)
8896 let (!trained, !bestInput) = foldl' (\(!n, !i) offset -> trainSlice rate n i shakespeare offset size) (net, io) xs
89− let results = take 100 $ generateParagraph trained bestInput oneHotMap oneHotDictionary ( S1D $ konst 0)
97+ results <- take 1000 <$> generateParagraph trained bestInput oneHotMap oneHotDictionary ( S1D $ konst 0)
9098 putStrLn ("TRAINING STEP WITH SIZE: " ++ show size)
9199 putStrLn (unAnnotateCapitals results)
92100 return (trained, bestInput)
93− ) (net0, i0) [10,10,15,15,20,20,25,25,30,30,35,35,40,40,50 :: Int]
101+ ) (net0, i0) [50,50,50,50,50,50,50,50,50,50,50,50,50,50,50 :: Int]
102+
103+ case savePath of
104+ Just saveFile -> lift . B.writeFile saveFile $ runPut (put (trained, bestInput))
105+ Nothing -> return ()
94106
95107generateParagraph :: forall layers shapes n a. (Last shapes ~ 'D1 n, Head shapes ~ 'D1 n, KnownNat n, Ord a)
96108 => RecurrentNetwork layers shapes
@@ -98,20 +110,23 @@ generateParagraph :: forall layers shapes n a. (Last shapes ~ 'D1 n, Head shapes
98110 -> M.Map a Int
99111 -> Vector a
100112 -> S ('D1 n)
101− -> [a]
102−generateParagraph n s hotmap hotdict i =
103− unfoldr go (s, i)
113+ -> IO [a]
114+generateParagraph n s hotmap hotdict =
115+ go s
104116 where
105− go (x, y) =
117+ go x y =
106118 do let (ns, o) = runRecurrent n x y
107− un <- unHot hotdict o
108− re <- makeHot hotmap un
109− Just (un, (ns, re))
119+ un <- sample 0.4 hotdict o
120+ Just re <- return $ makeHot hotmap un
121+ rest <- unsafeInterleaveIO $ go ns re
122+ return (un : rest)
110123
111124data ShakespeareOpts = ShakespeareOpts {
112125 trainingFile :: FilePath
113126 , iterations :: Int
114127 , rate :: LearningParameters
128+ , loadPath :: Maybe FilePath
129+ , savePath :: Maybe FilePath
115130 }
116131
117132shakespeare' :: Parser ShakespeareOpts
@@ -122,6 +137,8 @@ shakespeare' = ShakespeareOpts <$> argument str (metavar "TRAIN")
122137 <*> option auto (long "momentum" <> value 0.95)
123138 <*> option auto (long "l2" <> value 0.000001)
124139 )
140+ <*> optional (strOption (long "load"))
141+ <*> optional (strOption (long "save"))
125142
126143main :: IO ()
127144main = do
@@ -132,6 +149,11 @@ main = do
132149 Left err -> putStrLn err
133150
134151
152+netLoad :: FilePath -> IO (Shakespeare, Shakespearian)
153+netLoad modelPath = do
154+ modelData <- B.readFile modelPath
155+ either fail return $ runGet get modelData
156+
135157-- Replace capitals with an annotation and the lower case letter
136158-- http://fastml.com/one-weird-trick-for-training-char-rnns/
137159annotateCapitals :: String -> String