git.iko.soy/grenade
on master help sign in
Optimise Wengert tape for LSTM
a588151Huw Campbell
bench/bench-lstm.hs+1−40
compare with master
unifiedsplit
@@ -5,61 +5,22 @@ import Criterion.Main
5
5
6import Grenade
6import Grenade
7import Grenade.Recurrent
7import Grenade.Recurrent
8−import Grenade.Layers.Internal.Update
9−
10−import qualified Numeric.LinearAlgebra.Static as H
11
8
12main :: IO ()
9main :: IO ()
13main = do
10main = do
14− layer60 :: LSTM 40 60 <- createRandom
15− layer512 :: LSTM 40 512 <- createRandom
16 input40 :: S ('D1 40) <- randomOfShape
11 input40 :: S ('D1 40) <- randomOfShape
17− rec60 :: S ('D1 60) <- randomOfShape
18− rec512 :: S ('D1 512) <- randomOfShape
19 lstm :: RecNet <- randomRecurrent
12 lstm :: RecNet <- randomRecurrent
20
13
21− let upIn60 :: H.R 3600 = H.randomVector 1 H.Uniform * 2 - 1
22− let upIn512 :: H.R 262144 = H.randomVector 1 H.Uniform * 2 - 1
23−
24 defaultMain [
14 defaultMain [
25− bgroup "lstm" [ bench "forwards-60" $ nf (nfT3 . uncurry (testRun60 layer60)) (rec60, input40)
26− , bench "forwards-512" $ nf (nfT3 . uncurry (testRun512 layer512)) (rec512, input40)
27− , bench "backwards-60" $ nf (nfT3 . uncurry4 (testRun60' layer60)) (rec60, input40, rec60, rec60)
28− , bench "backwards-512" $ nf (nfT3 . uncurry4 (testRun512' layer512)) (rec512, input40, rec512, rec512)
29− ]
30− , bgroup "update" [ bench "matrix-60x60" $ nf (uncurry3 (descendVector 1 1 1)) (upIn60, upIn60, upIn60)
31− , bench "matrix-512x512" $ nf (uncurry3 (descendVector 1 1 1)) (upIn512, upIn512, upIn512)
32− ]
33− , bgroup "train" [ bench "one-time-step" $ whnf (nfT2 . trainRecurrent lp lstm 0) [(input40, Just input40)]
15+ bgroup "train" [ bench "one-time-step" $ whnf (nfT2 . trainRecurrent lp lstm 0) [(input40, Just input40)]
34 , bench "ten-time-steps" $ whnf (nfT2 . trainRecurrent lp lstm 0) $ replicate 10 (input40, Just input40)
16 , bench "ten-time-steps" $ whnf (nfT2 . trainRecurrent lp lstm 0) $ replicate 10 (input40, Just input40)
35 , bench "fifty-time-steps" $ whnf (nfT2 . trainRecurrent lp lstm 0) $ replicate 50 (input40, Just input40)
17 , bench "fifty-time-steps" $ whnf (nfT2 . trainRecurrent lp lstm 0) $ replicate 50 (input40, Just input40)
36 ]
18 ]
37 ]
19 ]
38
20
39−testRun60 :: LSTM 40 60 -> S ('D1 60) -> S ('D1 40) -> ((S ('D1 60), S ('D1 40)), S ('D1 60), S ('D1 60))
40−testRun60 = runRecurrentForwards
41−
42−testRun60' :: LSTM 40 60 -> S ('D1 60) -> S ('D1 40) -> S ('D1 60) -> S ('D1 60) -> (Gradient (LSTM 40 60), S ('D1 60), S ('D1 40))
43−testRun60' = curry . runRecurrentBackwards
44−
45−testRun512 :: LSTM 40 512 -> S ('D1 512) -> S ('D1 40) -> ((S ('D1 512), S ('D1 40)), S ('D1 512), S ('D1 512))
46−testRun512 = runRecurrentForwards
47−
48−testRun512' :: LSTM 40 512 -> S ('D1 512) -> S ('D1 40) -> S ('D1 512) -> S ('D1 512) -> (Gradient (LSTM 40 512), S ('D1 512), S ('D1 40))
49−testRun512' = curry . runRecurrentBackwards
50−
51−uncurry4 :: (t -> t1 -> t2 -> t3 -> t4) -> (t, t1, t2, t3) -> t4
52−uncurry4 f (a,b,c,d) = f a b c d
53−
54−uncurry3 :: (t -> t1 -> t2 -> t3) -> (t, t1, t2) -> t3
55−uncurry3 f (a,b,c) = f a b c
56−
57nfT2 :: (a, b) -> (a, b)
21nfT2 :: (a, b) -> (a, b)
58nfT2 (!a, !b) = (a, b)
22nfT2 (!a, !b) = (a, b)
59
23
60−nfT3 :: (a, b, c) -> (b, c)
61−nfT3 (!_, !b, !c) = (b, c)
62−
63
24
64type R = Recurrent
25type R = Recurrent
65type RecNet = RecurrentNetwork '[ R (LSTM 40 512), R (LSTM 512 40) ]
26type RecNet = RecurrentNetwork '[ R (LSTM 40 512), R (LSTM 512 40) ]
src/Grenade/Layers/Elu.hs+3−3
@@ -36,12 +36,12 @@ instance Serialize Elu where
36 get = return Elu
36 get = return Elu
37
37
38instance ( KnownNat i) => Layer Elu ('D1 i) ('D1 i) where
38instance ( KnownNat i) => Layer Elu ('D1 i) ('D1 i) where
39− type Tape Elu ('D1 i) ('D1 i) = S ('D1 i)
39+ type Tape Elu ('D1 i) ('D1 i) = LAS.R i
40
40
41− runForwards _ (S1D y) = (S1D y, S1D (elu y))
41+ runForwards _ (S1D y) = (y, S1D (elu y))
42 where
42 where
43 elu = LAS.dvmap (\a -> if a <= 0 then exp a - 1 else a)
43 elu = LAS.dvmap (\a -> if a <= 0 then exp a - 1 else a)
44− runBackwards _ (S1D y) (S1D dEdy) = ((), S1D (elu' y * dEdy))
44+ runBackwards _ y (S1D dEdy) = ((), S1D (elu' y * dEdy))
45 where
45 where
46 elu' = LAS.dvmap (\a -> if a <= 0 then exp a else 1)
46 elu' = LAS.dvmap (\a -> if a <= 0 then exp a else 1)
47
47
src/Grenade/Layers/FullyConnected.hs+3−3
@@ -46,12 +46,12 @@ instance (KnownNat i, KnownNat o) => UpdateLayer (FullyConnected i o) where
46 createRandom = randomFullyConnected
46 createRandom = randomFullyConnected
47
47
48instance (KnownNat i, KnownNat o) => Layer (FullyConnected i o) ('D1 i) ('D1 o) where
48instance (KnownNat i, KnownNat o) => Layer (FullyConnected i o) ('D1 i) ('D1 o) where
49− type Tape (FullyConnected i o) ('D1 i) ('D1 o) = S ('D1 i)
49+ type Tape (FullyConnected i o) ('D1 i) ('D1 o) = R i
50 -- Do a matrix vector multiplication and return the result.
50 -- Do a matrix vector multiplication and return the result.
51− runForwards (FullyConnected (FullyConnected' wB wN) _) (S1D v) = (S1D v, S1D (wB + wN #> v))
51+ runForwards (FullyConnected (FullyConnected' wB wN) _) (S1D v) = (v, S1D (wB + wN #> v))
52
52
53 -- Run a backpropogation step for a full connected layer.
53 -- Run a backpropogation step for a full connected layer.
54− runBackwards (FullyConnected (FullyConnected' _ wN) _) (S1D x) (S1D dEdy) =
54+ runBackwards (FullyConnected (FullyConnected' _ wN) _) x (S1D dEdy) =
55 let wB' = dEdy
55 let wB' = dEdy
56 mm' = dEdy `outer` x
56 mm' = dEdy `outer` x
57 -- calcluate derivatives for next step
57 -- calcluate derivatives for next step
src/Grenade/Layers/Logit.hs+15−9
@@ -34,18 +34,24 @@ instance UpdateLayer Logit where
34 createRandom = return Logit
34 createRandom = return Logit
35
35
36instance (a ~ b, SingI a) => Layer Logit a b where
36instance (a ~ b, SingI a) => Layer Logit a b where
37+ -- Wengert tape optimisation:
38+ --
39+ -- Derivative of the sigmoid function is
40+ -- d σ(x) / dx = σ(x) • (1 - σ(x))
41+ -- but we have already calculated σ(x) in
42+ -- the forward pass, so just store that
43+ -- and use it in the backwards pass.
37 type Tape Logit a b = S a
44 type Tape Logit a b = S a
38− runForwards _ a = (a, logistic a)
45+ runForwards _ a =
46+ let l = sigmoid a
47+ in (l, l)
39− runBackwards _ a g = ((), logistic' a * g)
48+ runBackwards _ l g =
49+ let sigmoid' = l * (1 - l)
50+ in ((), sigmoid' * g)
40
51
41instance Serialize Logit where
52instance Serialize Logit where
42 put _ = return ()
53 put _ = return ()
43 get = return Logit
54 get = return Logit
44
55
45−logistic :: Floating a => a -> a
56+sigmoid :: Floating a => a -> a
46−logistic x = 1 / (1 + exp (-x))
57+sigmoid x = 1 / (1 + exp (-x))
47−
48−logistic' :: Floating a => a -> a
49−logistic' x = logix * (1 - logix)
50− where
51− logix = logistic x
src/Grenade/Recurrent/Layers/LSTM.hs+16−30
@@ -127,37 +127,12 @@ instance (KnownNat i, KnownNat o) => RecurrentUpdateLayer (LSTM i o) where
127
127
128instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) where
128instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) where
129
129
130+ -- The tape stores essentially every variable we calculate,
131+ -- so we don't have to run any forwards component again.
130− type RecTape (LSTM i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
132+ 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)
131 -- Forward propagation for the LSTM layer.
133 -- Forward propagation for the LSTM layer.
132 -- The size of the cell state is also the size of the output.
134 -- The size of the cell state is also the size of the output.
133 runRecurrentForwards (LSTM (LSTMWeights {..}) _) (S1D cell) (S1D input) =
135 runRecurrentForwards (LSTM (LSTMWeights {..}) _) (S1D cell) (S1D input) =
134− let -- Forget state vector
135− f_t = sigmoid $ lstmBf + lstmWf #> input + lstmUf #> cell
136− -- Input state vector
137− i_t = sigmoid $ lstmBi + lstmWi #> input + lstmUi #> cell
138− -- Output state vector
139− o_t = sigmoid $ lstmBo + lstmWo #> input + lstmUo #> cell
140− -- Cell input state vector
141− c_x = tanh $ lstmBc + lstmWc #> input
142− -- Cell state
143− c_t = f_t * cell + i_t * c_x
144− -- Output (it's sometimes recommended to use tanh c_t)
145− h_t = o_t * c_t
146− in ((S1D cell, S1D input), S1D c_t, S1D h_t)
147−
148− -- Run a backpropogation step for an LSTM layer.
149− -- We're doing all the derivatives by hand here, so one should
150− -- be extra careful when changing this.
151− --
152− -- There's a test version using the AD library without hmatrix in the test
153− -- 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
136 let -- Forget state vector
162 f_s = lstmBf + lstmWf #> input + lstmUf #> cell
137 f_s = lstmBf + lstmWf #> input + lstmUf #> cell
163 f_t = sigmoid f_s
138 f_t = sigmoid f_s
@@ -172,8 +147,18 @@ instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) w
172 c_x = tanh c_s
147 c_x = tanh c_s
173 -- Cell state
148 -- Cell state
174 c_t = f_t * cell + i_t * c_x
149 c_t = f_t * cell + i_t * c_x
150+ -- Output (it's sometimes recommended to use tanh c_t)
151+ h_t = o_t * c_t
152+ 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)
175
153
154+ -- Run a backpropogation step for an LSTM layer.
155+ -- We're doing all the derivatives by hand here, so one should
156+ -- be extra careful when changing this.
157+ --
158+ -- There's a test version using the AD library without hmatrix in the test
159+ -- suite. These should match always.
160+ runRecurrentBackwards (LSTM (LSTMWeights {..}) _) (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') =
176− -- Reverse Mode AD Derivitives
161+ let -- Reverse Mode AD Derivitives
177 c_t' = h_t' * o_t + cellGrad
162 c_t' = h_t' * o_t + cellGrad
178
163
179 f_t' = c_t' * cell
164 f_t' = c_t' * cell
@@ -235,7 +220,8 @@ randomLSTM = do
235
220
236-- | Maths
221-- | Maths
237--
222--
238−-- TODO: move to not here
223+-- TODO: Move to not here
224+-- Optimise backwards derivative
239sigmoid :: Floating a => a -> a
225sigmoid :: Floating a => a -> a
240sigmoid x = 1 / (1 + exp (-x))
226sigmoid x = 1 / (1 + exp (-x))
241
227
test/Test/Grenade/Recurrent/Layers/LSTM.hs+9−3
@@ -65,7 +65,9 @@ prop_lstm_reference_backwards =
65 input :: S.R 3 <- forAll randomVector
65 input :: S.R 3 <- forAll randomVector
66 cell :: S.R 2 <- forAll randomVector
66 cell :: S.R 2 <- forAll randomVector
67 net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
67 net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
68− let actualBacks = runRecurrentBackwards net (S1D cell, S1D input) (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
68+ let (tape, _ :: S ('D1 2), _ :: S ('D1 2))
69+ = runRecurrentForwards net (S1D cell) (S1D input)
70+ actualBacks = runRecurrentBackwards net tape (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
69 case actualBacks of
71 case actualBacks of
70 (actualGradients, _, _ :: S ('D1 3)) ->
72 (actualGradients, _, _ :: S ('D1 3)) ->
71 let refNet = Reference.lstmToReference lstmWeights
73 let refNet = Reference.lstmToReference lstmWeights
@@ -79,7 +81,9 @@ prop_lstm_reference_backwards_input =
79 input :: S.R 3 <- forAll randomVector
81 input :: S.R 3 <- forAll randomVector
80 cell :: S.R 2 <- forAll randomVector
82 cell :: S.R 2 <- forAll randomVector
81 net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
83 net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
82− let actualBacks = runRecurrentBackwards net (S1D cell, S1D input) (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
84+ let (tape, _ :: S ('D1 2), _ :: S ('D1 2))
85+ = runRecurrentForwards net (S1D cell) (S1D input)
86+ actualBacks = runRecurrentBackwards net tape (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
83 case actualBacks of
87 case actualBacks of
84 (_, _, S1D actualGradients :: S ('D1 3)) ->
88 (_, _, S1D actualGradients :: S ('D1 3)) ->
85 let refNet = Reference.lstmToReference lstmWeights
89 let refNet = Reference.lstmToReference lstmWeights
@@ -93,7 +97,9 @@ prop_lstm_reference_backwards_cell =
93 input :: S.R 3 <- forAll randomVector
97 input :: S.R 3 <- forAll randomVector
94 cell :: S.R 2 <- forAll randomVector
98 cell :: S.R 2 <- forAll randomVector
95 net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
99 net@(LSTM lstmWeights _) :: LSTM 3 2 <- forAll genLSTM
96− let actualBacks = runRecurrentBackwards net (S1D cell, S1D input) (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
100+ let (tape, _ :: S ('D1 2), _ :: S ('D1 2))
101+ = runRecurrentForwards net (S1D cell) (S1D input)
102+ actualBacks = runRecurrentBackwards net tape (S1D (S.konst 1) :: S ('D1 2)) (S1D (S.konst 1) :: S ('D1 2))
97 case actualBacks of
103 case actualBacks of
98 (_, S1D actualGradients, _ :: S ('D1 3)) ->
104 (_, S1D actualGradients, _ :: S ('D1 3)) ->
99 let refNet = Reference.lstmToReference lstmWeights
105 let refNet = Reference.lstmToReference lstmWeights