git.iko.soy/grenade
on master help sign in
grenade/bench/bench-lstm.hs
30 lines 994 B history raw
1{-# LANGUAGE DataKinds #-}
2{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ScopedTypeVariables #-}
4import Criterion.Main
6import Grenade
7import Grenade.Recurrent
9main :: IO ()
10main = do
11 input40 :: S ('D1 40) <- randomOfShape
12 lstm :: RecNet <- randomRecurrent
14 defaultMain [
15 bgroup "train" [ bench "one-time-step" $ whnf (nfT2 . trainRecurrent lp lstm 0) [(input40, Just input40)]
16 , bench "ten-time-steps" $ whnf (nfT2 . trainRecurrent lp lstm 0) $ replicate 10 (input40, Just input40)
17 , bench "fifty-time-steps" $ whnf (nfT2 . trainRecurrent lp lstm 0) $ replicate 50 (input40, Just input40)
18 ]
19 ]
21nfT2 :: (a, b) -> (a, b)
22nfT2 (!a, !b) = (a, b)
25type R = Recurrent
26type RecNet = RecurrentNetwork '[ R (LSTM 40 512), R (LSTM 512 40) ]
27 '[ 'D1 40, 'D1 512, 'D1 40 ]
29lp :: LearningParameters
30lp = LearningParameters 0.1 0 0