@@ -1,8 +1,8 @@
1−{-# LANGUAGE CPP #-}
21{-# LANGUAGE BangPatterns #-}
32{-# LANGUAGE DataKinds #-}
43{-# LANGUAGE GADTs #-}
54{-# 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.Kind (Type)
24+import Data.Singletons.TypeLits
2525
2626import qualified Numeric.LinearAlgebra as LA
2727import Numeric.LinearAlgebra.Static
@@ -29,7 +29,6 @@ import Numeric.LinearAlgebra.Static
2929import Grenade.Core
3030import Grenade.Recurrent.Core
3131import Grenade.Layers.Internal.Update
32−import GHC.TypeLits
3332
3433
3534-- | Long Short Term Memory Recurrent unit
@@ -37,14 +36,14 @@ import GHC.TypeLits
3736-- This is a Peephole formulation, so the recurrent shape is
3837-- just the cell state, the previous output is not held or used
3938-- at all.
40−data LSTM :: Nat -> Nat -> Type where
39+data LSTM :: Nat -> Nat -> * where
4140 LSTM :: ( KnownNat input
4241 , KnownNat output
4342 ) => !(LSTMWeights input output) -- Weights
4443 -> !(LSTMWeights input output) -- Momentums
4544 -> LSTM input output
4645
47−data LSTMWeights :: Nat -> Nat -> Type where
46+data LSTMWeights :: Nat -> Nat -> * where
4847 LSTMWeights :: ( KnownNat input
4948 , KnownNat output
5049 ) => {
@@ -71,7 +70,7 @@ instance (KnownNat i, KnownNat o) => UpdateLayer (LSTM i o) where
7170
7271 -- Run the update function for each group matrix/vector of weights, momentums and gradients.
7372 -- Hmm, maybe the function should be used instead of passing in the learning parameters.
74− runUpdate lp (LSTM w m) g =
73+ runUpdate LearningParameters {..} (LSTM w m) g =
7574 let (wf, wf') = u lstmWf w m g
7675 (uf, uf') = u lstmUf w m g
7776 (bf, bf') = v lstmBf w m g
@@ -88,11 +87,11 @@ instance (KnownNat i, KnownNat o) => UpdateLayer (LSTM i o) where
8887 -- Utility function for updating with the momentum, gradients, and weights.
8988 u :: forall x ix out. (KnownNat ix, KnownNat out) => (x -> (L out ix)) -> x -> x -> x -> ((L out ix), (L out ix))
9089 u e (e -> weights) (e -> momentum) (e -> gradient) =
91− descendMatrix (learningRate lp) (learningMomentum lp) (learningRegulariser lp) weights gradient momentum
90+ descendMatrix learningRate learningMomentum learningRegulariser weights gradient momentum
9291
9392 v :: forall x ix. (KnownNat ix) => (x -> (R ix)) -> x -> x -> x -> ((R ix), (R ix))
9493 v e (e -> weights) (e -> momentum) (e -> gradient) =
95− descendVector (learningRate lp) (learningMomentum lp) (learningRegulariser lp) weights gradient momentum
94+ descendVector learningRate learningMomentum learningRegulariser weights gradient momentum
9695
9796 -- There's a lot of updates here, so to try and minimise the number of data copies
9897 -- we'll create a mutable bucket for each.
@@ -128,29 +127,23 @@ instance (KnownNat i, KnownNat o) => RecurrentUpdateLayer (LSTM i o) where
128127
129128instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) where
130129
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)
130+ type RecTape (LSTM i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
134131 -- Forward propagation for the LSTM layer.
135132 -- The size of the cell state is also the size of the output.
136− runRecurrentForwards (LSTM lw _) (S1D cell) (S1D input) =
133+ runRecurrentForwards (LSTM (LSTMWeights {..}) _) (S1D cell) (S1D input) =
137134 let -- Forget state vector
138− f_s = lstmBf lw + lstmWf lw #> input + lstmUf lw #> cell
139− f_t = sigmoid f_s
135+ f_t = sigmoid $ lstmBf + lstmWf #> input + lstmUf #> cell
140136 -- Input state vector
141− i_s = lstmBi lw + lstmWi lw #> input + lstmUi lw #> cell
142− i_t = sigmoid i_s
137+ i_t = sigmoid $ lstmBi + lstmWi #> input + lstmUi #> cell
143138 -- Output state vector
144− o_s = lstmBo lw + lstmWo lw #> input + lstmUo lw #> cell
145− o_t = sigmoid o_s
139+ o_t = sigmoid $ lstmBo + lstmWo #> input + lstmUo #> cell
146140 -- Cell input state vector
147− c_s = lstmBc lw + lstmWc lw #> input
148− c_x = tanh c_s
141+ c_x = tanh $ lstmBc + lstmWc #> input
149142 -- Cell state
150143 c_t = f_t * cell + i_t * c_x
151144 -- Output (it's sometimes recommended to use tanh c_t)
152145 h_t = o_t * c_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)
146+ in ((S1D cell, S1D input), S1D c_t, S1D h_t)
154147
155148 -- Run a backpropogation step for an LSTM layer.
156149 -- We're doing all the derivatives by hand here, so one should
@@ -158,8 +151,29 @@ instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) w
158151 --
159152 -- There's a test version using the AD library without hmatrix in the test
160153 -- suite. These should match always.
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
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
163177 c_t' = h_t' * o_t + cellGrad
164178
165179 f_t' = c_t' * cell
@@ -175,8 +189,8 @@ instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) w
175189 c_s' = tanh' c_s * c_x'
176190
177191 -- The derivatives to pass sideways (recurrent) and downwards
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'
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'
180194
181195 -- Calculate the gradient Matricies for the input
182196 lstmWf' = f_s' `outer` input
@@ -221,8 +235,7 @@ randomLSTM = do
221235
222236-- | Maths
223237--
224−-- TODO: Move to not here
225−-- Optimise backwards derivative
238+-- TODO: move to not here
226239sigmoid :: Floating a => a -> a
227240sigmoid x = 1 / (1 + exp (-x))
228241
@@ -235,34 +248,44 @@ tanh' :: (Floating a) => a -> a
235248tanh' t = 1 - s ^ (2 :: Int) where s = tanh t
236249
237250instance (KnownNat i, KnownNat o) => Serialize (LSTM i o) where
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
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
255268
256269 get = do
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
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