@@ -1,8 +1,8 @@
1+{-# LANGUAGE CPP #-}
12{-# LANGUAGE BangPatterns #-}
23{-# LANGUAGE DataKinds #-}
34{-# LANGUAGE GADTs #-}
45{-# LANGUAGE RankNTypes #-}
5−{-# LANGUAGE RecordWildCards #-}
66{-# LANGUAGE TypeOperators #-}
77{-# LANGUAGE TypeFamilies #-}
88{-# LANGUAGE MultiParamTypeClasses #-}
@@ -21,7 +21,7 @@ import Control.Monad.Random ( MonadRandom, getRandom )
2121-- import Data.List ( foldl1' )
2222import Data.Proxy
2323import Data.Serialize
24−import Data.Singletons.TypeLits
24+import Data.Kind (Type)
2525
2626import qualified Numeric.LinearAlgebra as LA
2727import Numeric.LinearAlgebra.Static
@@ -29,6 +29,7 @@ import Numeric.LinearAlgebra.Static
2929import Grenade.Core
3030import Grenade.Recurrent.Core
3131import Grenade.Layers.Internal.Update
32+import GHC.TypeLits
3233
3334
3435-- | Long Short Term Memory Recurrent unit
@@ -36,14 +37,14 @@ import Grenade.Layers.Internal.Update
3637-- This is a Peephole formulation, so the recurrent shape is
3738-- just the cell state, the previous output is not held or used
3839-- at all.
39−data LSTM :: Nat -> Nat -> * where
40+data LSTM :: Nat -> Nat -> Type where
4041 LSTM :: ( KnownNat input
4142 , KnownNat output
4243 ) => !(LSTMWeights input output) -- Weights
4344 -> !(LSTMWeights input output) -- Momentums
4445 -> LSTM input output
4546
46−data LSTMWeights :: Nat -> Nat -> * where
47+data LSTMWeights :: Nat -> Nat -> Type where
4748 LSTMWeights :: ( KnownNat input
4849 , KnownNat output
4950 ) => {
@@ -70,7 +71,7 @@ instance (KnownNat i, KnownNat o) => UpdateLayer (LSTM i o) where
7071
7172 -- Run the update function for each group matrix/vector of weights, momentums and gradients.
7273 -- Hmm, maybe the function should be used instead of passing in the learning parameters.
73− runUpdate LearningParameters {..} (LSTM w m) g =
74+ runUpdate lp (LSTM w m) g =
7475 let (wf, wf') = u lstmWf w m g
7576 (uf, uf') = u lstmUf w m g
7677 (bf, bf') = v lstmBf w m g
@@ -87,11 +88,11 @@ instance (KnownNat i, KnownNat o) => UpdateLayer (LSTM i o) where
8788 -- Utility function for updating with the momentum, gradients, and weights.
8889 u :: forall x ix out. (KnownNat ix, KnownNat out) => (x -> (L out ix)) -> x -> x -> x -> ((L out ix), (L out ix))
8990 u e (e -> weights) (e -> momentum) (e -> gradient) =
90− descendMatrix learningRate learningMomentum learningRegulariser weights gradient momentum
91+ descendMatrix (learningRate lp) (learningMomentum lp) (learningRegulariser lp) weights gradient momentum
9192
9293 v :: forall x ix. (KnownNat ix) => (x -> (R ix)) -> x -> x -> x -> ((R ix), (R ix))
9394 v e (e -> weights) (e -> momentum) (e -> gradient) =
94− descendVector learningRate learningMomentum learningRegulariser weights gradient momentum
95+ descendVector (learningRate lp) (learningMomentum lp) (learningRegulariser lp) weights gradient momentum
9596
9697 -- There's a lot of updates here, so to try and minimise the number of data copies
9798 -- we'll create a mutable bucket for each.
@@ -127,23 +128,29 @@ instance (KnownNat i, KnownNat o) => RecurrentUpdateLayer (LSTM i o) where
127128
128129instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) where
129130
130− type RecTape (LSTM i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
131+ -- The tape stores essentially every variable we calculate,
132+ -- so we don't have to run any forwards component again.
133+ type RecTape (LSTM i o) ('D1 i) ('D1 o) = (R o, R i, R o, R o, R o, R o, R o, R o, R o, R o, R o)
131134 -- Forward propagation for the LSTM layer.
132135 -- The size of the cell state is also the size of the output.
133− runRecurrentForwards (LSTM (LSTMWeights {..}) _) (S1D cell) (S1D input) =
136+ runRecurrentForwards (LSTM lw _) (S1D cell) (S1D input) =
134137 let -- Forget state vector
135− f_t = sigmoid $ lstmBf + lstmWf #> input + lstmUf #> cell
138+ f_s = lstmBf lw + lstmWf lw #> input + lstmUf lw #> cell
139+ f_t = sigmoid f_s
136140 -- Input state vector
137− i_t = sigmoid $ lstmBi + lstmWi #> input + lstmUi #> cell
141+ i_s = lstmBi lw + lstmWi lw #> input + lstmUi lw #> cell
142+ i_t = sigmoid i_s
138143 -- Output state vector
139− o_t = sigmoid $ lstmBo + lstmWo #> input + lstmUo #> cell
144+ o_s = lstmBo lw + lstmWo lw #> input + lstmUo lw #> cell
145+ o_t = sigmoid o_s
140146 -- Cell input state vector
141− c_x = tanh $ lstmBc + lstmWc #> input
147+ c_s = lstmBc lw + lstmWc lw #> input
148+ c_x = tanh c_s
142149 -- Cell state
143150 c_t = f_t * cell + i_t * c_x
144151 -- Output (it's sometimes recommended to use tanh c_t)
145152 h_t = o_t * c_t
146− in ((S1D cell, S1D input), S1D c_t, S1D h_t)
153+ in ((cell, input, f_s, f_t, i_s, i_t, o_s, o_t, c_s, c_x, c_t), S1D c_t, S1D h_t)
147154
148155 -- Run a backpropogation step for an LSTM layer.
149156 -- We're doing all the derivatives by hand here, so one should
@@ -151,29 +158,8 @@ instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) w
151158 --
152159 -- There's a test version using the AD library without hmatrix in the test
153160 -- suite. These should match always.
154− runRecurrentBackwards (LSTM (LSTMWeights {..}) _) (S1D cell, S1D input) (S1D cellGrad) (S1D h_t') =
155− -- We're not keeping the Wengert tape during the forward pass,
156− -- so we're duplicating some work here.
157− --
158− -- If I was being generous, I'd call it checkpointing.
159− --
160− -- Maybe think about better ways to store some intermediate states.
161− let -- Forget state vector
162− f_s = lstmBf + lstmWf #> input + lstmUf #> cell
163− f_t = sigmoid f_s
164− -- Input state vector
165− i_s = lstmBi + lstmWi #> input + lstmUi #> cell
166− i_t = sigmoid i_s
167− -- Output state vector
168− o_s = lstmBo + lstmWo #> input + lstmUo #> cell
169− o_t = sigmoid o_s
170− -- Cell input state vector
171− c_s = lstmBc + lstmWc #> input
172− c_x = tanh c_s
173− -- Cell state
174− c_t = f_t * cell + i_t * c_x
175−
176− -- Reverse Mode AD Derivitives
161+ runRecurrentBackwards (LSTM lw _) (cell, input, f_s, f_t, i_s, i_t, o_s, o_t, c_s, c_x, c_t) (S1D cellGrad) (S1D h_t') =
162+ let -- Reverse Mode AD Derivitives
177163 c_t' = h_t' * o_t + cellGrad
178164
179165 f_t' = c_t' * cell
@@ -189,8 +175,8 @@ instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) w
189175 c_s' = tanh' c_s * c_x'
190176
191177 -- The derivatives to pass sideways (recurrent) and downwards
192− cell' = tr lstmUf #> f_s' + tr lstmUo #> o_s' + tr lstmUi #> i_s' + c_t' * f_t
193− input' = tr lstmWf #> f_s' + tr lstmWo #> o_s' + tr lstmWi #> i_s' + tr lstmWc #> c_s'
178+ cell' = tr (lstmUf lw) #> f_s' + tr (lstmUo lw) #> o_s' + tr (lstmUi lw) #> i_s' + c_t' * f_t
179+ input' = tr (lstmWf lw) #> f_s' + tr (lstmWo lw) #> o_s' + tr (lstmWi lw) #> i_s' + tr (lstmWc lw) #> c_s'
194180
195181 -- Calculate the gradient Matricies for the input
196182 lstmWf' = f_s' `outer` input
@@ -235,7 +221,8 @@ randomLSTM = do
235221
236222-- | Maths
237223--
238−-- TODO: move to not here
224+-- TODO: Move to not here
225+-- Optimise backwards derivative
239226sigmoid :: Floating a => a -> a
240227sigmoid x = 1 / (1 + exp (-x))
241228
@@ -248,44 +235,34 @@ tanh' :: (Floating a) => a -> a
248235tanh' t = 1 - s ^ (2 :: Int) where s = tanh t
249236
250237instance (KnownNat i, KnownNat o) => Serialize (LSTM i o) where
251− put (LSTM LSTMWeights {..} _) = do
252− u lstmWf
253− u lstmUf
254− v lstmBf
255− u lstmWi
256− u lstmUi
257− v lstmBi
258− u lstmWo
259− u lstmUo
260− v lstmBo
261− u lstmWc
262− v lstmBc
263− where
264− u :: forall a b. (KnownNat a, KnownNat b) => Putter (L b a)
265− u = putListOf put . LA.toList . LA.flatten . extract
266− v :: forall a. (KnownNat a) => Putter (R a)
267− v = putListOf put . LA.toList . extract
238+ put (LSTM lw _) = do
239+ u (lstmWf lw)
240+ u (lstmUf lw)
241+ v (lstmBf lw)
242+ u (lstmWi lw)
243+ u (lstmUi lw)
244+ v (lstmBi lw)
245+ u (lstmWo lw)
246+ u (lstmUo lw)
247+ v (lstmBo lw)
248+ u (lstmWc lw)
249+ v (lstmBc lw)
250+ where
251+ u :: forall a b. (KnownNat a, KnownNat b) => Putter (L b a)
252+ u = putListOf put . LA.toList . LA.flatten . extract
253+ v :: forall a. (KnownNat a) => Putter (R a)
254+ v = putListOf put . LA.toList . extract
268255
269256 get = do
270− lstmWf <- u
271− lstmUf <- u
272− lstmBf <- v
273− lstmWi <- u
274− lstmUi <- u
275− lstmBi <- v
276− lstmWo <- u
277− lstmUo <- u
278− lstmBo <- v
279− lstmWc <- u
280− lstmBc <- v
281− return $ LSTM (LSTMWeights {..}) (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)
282− where
283− u :: forall a b. (KnownNat a, KnownNat b) => Get (L b a)
284− u = let f = fromIntegral $ natVal (Proxy :: Proxy a)
285− in maybe (fail "Vector of incorrect size") return . create . LA.reshape f . LA.fromList =<< getListOf get
286− v :: forall a. (KnownNat a) => Get (R a)
287− v = maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
288−
289− w0 = konst 0
290− u0 = konst 0
291− v0 = konst 0
257+ w <- LSTMWeights <$> u <*> u <*> v <*> u <*> u <*> v <*> u <*> u <*> v <*> u <*> v
258+ return $ LSTM w (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)
259+ where
260+ u :: forall a b. (KnownNat a, KnownNat b) => Get (L b a)
261+ u = let f = fromIntegral $ natVal (Proxy :: Proxy a)
262+ in maybe (fail "Vector of incorrect size") return . create . LA.reshape f . LA.fromList =<< getListOf get
263+ v :: forall a. (KnownNat a) => Get (R a)
264+ v = maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
265+
266+ w0 = konst 0
267+ u0 = konst 0
268+ v0 = konst 0