git.iko.soy/grenade
on master help sign in
Optimise Wengert tape for LSTM
commita5881518bf69d6033e0ccfae21f7b3ea509305a3 authorHuw Campbell <huw.campbell@ambiata.com> dateSat 16 Dec 2017, 21:21 (9 years ago) parent4332d62
6 files changed+47−88
compare with master
unifiedsplit
bench/bench-lstm.hs+1−40 src/Grenade/Layers/Elu.hs+3−3 src/Grenade/Layers/FullyConnected.hs+3−3 src/Grenade/Layers/Logit.hs+15−9 src/Grenade/Recurrent/Layers/LSTM.hs+16−30 test/Test/Grenade/Recurrent/Layers/LSTM.hs+9−3
bench/bench-lstm.hs+1−40
@@ -5,61 +5,22 @@ import Criterion.Main
55
66import Grenade
77import Grenade.Recurrent
8−import Grenade.Layers.Internal.Update
9−
10−import qualified Numeric.LinearAlgebra.Static as H
118
129main :: IO ()
1310main = do
14− layer60 :: LSTM 40 60 <- createRandom
15− layer512 :: LSTM 40 512 <- createRandom
1611 input40 :: S ('D1 40) <- randomOfShape
17− rec60 :: S ('D1 60) <- randomOfShape
18− rec512 :: S ('D1 512) <- randomOfShape
1912 lstm :: RecNet <- randomRecurrent
2013
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−
2414 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)]
3416 , bench "ten-time-steps" $ whnf (nfT2 . trainRecurrent lp lstm 0) $ replicate 10 (input40, Just input40)
3517 , bench "fifty-time-steps" $ whnf (nfT2 . trainRecurrent lp lstm 0) $ replicate 50 (input40, Just input40)
3618 ]
3719 ]
3820
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−
5721nfT2 :: (a, b) -> (a, b)
5822nfT2 (!a, !b) = (a, b)
5923
60−nfT3 :: (a, b, c) -> (b, c)
61−nfT3 (!_, !b, !c) = (b, c)
62−
6324
6425type R = Recurrent
6526type 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
3636 get = return Elu
3737
3838instance ( 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
4040
41− runForwards _ (S1D y) = (S1D y, S1D (elu y))
41+ runForwards _ (S1D y) = (y, S1D (elu y))
4242 where
4343 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))
4545 where
4646 elu' = LAS.dvmap (\a -> if a <= 0 then exp a else 1)
4747
src/Grenade/Layers/FullyConnected.hs+3−3
@@ -46,12 +46,12 @@ instance (KnownNat i, KnownNat o) => UpdateLayer (FullyConnected i o) where
4646 createRandom = randomFullyConnected
4747
4848instance (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
5050 -- 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))
5252
5353 -- 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) =
5555 let wB' = dEdy
5656 mm' = dEdy `outer` x
5757 -- calcluate derivatives for next step
src/Grenade/Layers/Logit.hs+15−9
@@ -34,18 +34,24 @@ instance UpdateLayer Logit where
3434 createRandom = return Logit
3535
3636instance (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.
3744 type Tape Logit a b = S a
38− runForwards _ a = (a, logistic a)
39− runBackwards _ a g = ((), logistic' a * g)
45+ runForwards _ a =
46+ let l = sigmoid a
47+ in (l, l)
48+ runBackwards _ l g =
49+ let sigmoid' = l * (1 - l)
50+ in ((), sigmoid' * g)
4051
4152instance Serialize Logit where
4253 put _ = return ()
4354 get = return Logit
4455
45−logistic :: Floating a => a -> a
46−logistic x = 1 / (1 + exp (-x))
47−
48−logistic' :: Floating a => a -> a
49−logistic' x = logix * (1 - logix)
50− where
51− logix = logistic x
56+sigmoid :: Floating a => a -> a
57+sigmoid x = 1 / (1 + exp (-x))
src/Grenade/Recurrent/Layers/LSTM.hs+16−30
@@ -127,37 +127,12 @@ instance (KnownNat i, KnownNat o) => RecurrentUpdateLayer (LSTM i o) where
127127
128128instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) where
129129
130− type RecTape (LSTM i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
130+ -- The tape stores essentially every variable we calculate,
131+ -- so we don't have to run any forwards component again.
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)
131133 -- Forward propagation for the LSTM layer.
132134 -- The size of the cell state is also the size of the output.
133135 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.
161136 let -- Forget state vector
162137 f_s = lstmBf + lstmWf #> input + lstmUf #> cell
163138 f_t = sigmoid f_s
@@ -172,8 +147,18 @@ instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) w
172147 c_x = tanh c_s
173148 -- Cell state
174149 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)
175153
176− -- Reverse Mode AD Derivitives
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') =
161+ let -- Reverse Mode AD Derivitives
177162 c_t' = h_t' * o_t + cellGrad
178163
179164 f_t' = c_t' * cell
@@ -235,7 +220,8 @@ randomLSTM = do
235220
236221-- | Maths
237222--
238−-- TODO: move to not here
223+-- TODO: Move to not here
224+-- Optimise backwards derivative
239225sigmoid :: Floating a => a -> a
240226sigmoid x = 1 / (1 + exp (-x))
241227
test/Test/Grenade/Recurrent/Layers/LSTM.hs+9−3
@@ -65,7 +65,9 @@ prop_lstm_reference_backwards =
6565 input :: S.R 3 <- forAll randomVector
6666 cell :: S.R 2 <- forAll randomVector
6767 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))
6971 case actualBacks of
7072 (actualGradients, _, _ :: S ('D1 3)) ->
7173 let refNet = Reference.lstmToReference lstmWeights
@@ -79,7 +81,9 @@ prop_lstm_reference_backwards_input =
7981 input :: S.R 3 <- forAll randomVector
8082 cell :: S.R 2 <- forAll randomVector
8183 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))
8387 case actualBacks of
8488 (_, _, S1D actualGradients :: S ('D1 3)) ->
8589 let refNet = Reference.lstmToReference lstmWeights
@@ -93,7 +97,9 @@ prop_lstm_reference_backwards_cell =
9397 input :: S.R 3 <- forAll randomVector
9498 cell :: S.R 2 <- forAll randomVector
9599 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))
97103 case actualBacks of
98104 (_, S1D actualGradients, _ :: S ('D1 3)) ->
99105 let refNet = Reference.lstmToReference lstmWeights