@@ -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 40), F (FullyConnected 40 40), F Logit]
54− '[ 'D1 40, '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 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,36 +86,50 @@ 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
87− xs <- take (iterations `div` 15) <$> getRandomRs (0, length shakespeare - size - 1)
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
95+ xs <- take (iterations `div` 10) <$> 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 temperature 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) $ replicate 10 sequenceSize
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
97109 -> RecurrentInputs layers
110+ -> Double
98111 -> M.Map a Int
99112 -> Vector a
100113 -> S ('D1 n)
101− -> [a]
102−generateParagraph n s hotmap hotdict i =
103− unfoldr go (s, i)
114+ -> IO [a]
115+generateParagraph n s temperature hotmap hotdict =
116+ go s
104117 where
105− go (x, y) =
118+ go x y =
106119 do let (ns, o) = runRecurrent n x y
107− un <- unHot hotdict o
108− re <- makeHot hotmap un
109− Just (un, (ns, re))
120+ un <- sample temperature hotdict o
121+ Just re <- return $ makeHot hotmap un
122+ rest <- unsafeInterleaveIO $ go ns re
123+ return (un : rest)
110124
111125data ShakespeareOpts = ShakespeareOpts {
112126 trainingFile :: FilePath
113127 , iterations :: Int
114128 , rate :: LearningParameters
129+ , sequenceSize :: Int
130+ , temperature :: Double
131+ , loadPath :: Maybe FilePath
132+ , savePath :: Maybe FilePath
115133 }
116134
117135shakespeare' :: Parser ShakespeareOpts
@@ -122,6 +140,10 @@ shakespeare' = ShakespeareOpts <$> argument str (metavar "TRAIN")
122140 <*> option auto (long "momentum" <> value 0.95)
123141 <*> option auto (long "l2" <> value 0.000001)
124142 )
143+ <*> option auto (long "sequence-length" <> short 's' <> value 50)
144+ <*> option auto (long "temperature" <> short 't' <> value 0.4)
145+ <*> optional (strOption (long "load"))
146+ <*> optional (strOption (long "save"))
125147
126148main :: IO ()
127149main = do
@@ -132,6 +154,11 @@ main = do
132154 Left err -> putStrLn err
133155
134156
157+netLoad :: FilePath -> IO (Shakespeare, Shakespearian)
158+netLoad modelPath = do
159+ modelData <- B.readFile modelPath
160+ either fail return $ runGet get modelData
161+
135162-- Replace capitals with an annotation and the lower case letter
136163-- http://fastml.com/one-weird-trick-for-training-char-rnns/
137164annotateCapitals :: String -> String