packages feed

mxnet-nn 0.0.1.1 → 0.0.1.2

raw patch · 8 files changed

+352/−94 lines, 8 filesdep ~mtlnew-component:exe:lenet

Dependency ranges changed: mtl

Files

examples/mnist/Dataset.hs view
@@ -20,10 +20,7 @@ type SymbolF = Symbol Float type ArrayF  = NDArray Float -device :: Context-device = contextCPU--type StreamProc a b m = Stream (Of a) m () -> Stream (Of b) m ()+type StreamProc a b m = Stream (Of a) m Int -> Stream (Of b) m Int  mappedOf :: Monad m => (a -> m b) -> StreamProc a b m -- mappedOf f = S.sequence . maps (first f)@@ -32,29 +29,35 @@ cImageToNDArray :: MonadIO m => StreamProc (Batched Image) ArrayF m cImageToNDArray = mappedOf $ \dat -> liftIO $ do   let sz = size dat-  makeNDArray [sz, 28, 28] device $ SV.concat $ NV.toList $ _batch dat+  makeNDArray [sz, 1, 28, 28] contextCPU $ SV.concat $ NV.toList $ _batch dat +cLabelToNDArray :: MonadIO m => StreamProc (Batched Label) ArrayF m+cLabelToNDArray = mappedOf $ \dat -> liftIO $ do+  let sz = size dat+  makeNDArray [sz] contextCPU (NV.convert $ NV.map fromIntegral $ _batch dat) :: IO ArrayF+ cLabelToOnehotNDArray :: MonadIO m => StreamProc (Batched Label) ArrayF m cLabelToOnehotNDArray = mappedOf $ \dat -> liftIO $ do   let sz = size dat-  a <- array [sz] (NV.convert $ NV.map fromIntegral $ _batch dat) :: IO ArrayF+  a <- makeNDArray [sz] contextCPU (NV.convert $ NV.map fromIntegral $ _batch dat) :: IO ArrayF   b <- MXI.one_hot (A.getHandle a) 10 (add @"on_value" 1.0 $ add @"off_value" 0.0 nil)   reshape (A.NDArray b) [sz, 10]  cBatchN :: MonadIO m => Int -> StreamProc a (Batched a) m-cBatchN n = mapped toBatch . chunksOf n+cBatchN n s = div' n <$> (mapped toBatch $ chunksOf n s)   where     toBatch seg = first (Batched . NV.fromList) <$> S.toList seg+    div' n t = let (r, m) = divMod t n in if m > 0 then r+1 else r   -trainingData :: MonadResource m => Stream (Of (ArrayF, ArrayF)) m ()+trainingData :: MonadResource m => Stream (Of (ArrayF, ArrayF)) m Int trainingData = S.zip     (sourceImages "examples/data/train-images-idx3-ubyte" & cBatchN 32 & cImageToNDArray      )-    (sourceLabels "examples/data/train-labels-idx1-ubyte" & cBatchN 32 & cLabelToOnehotNDArray)+    (sourceLabels "examples/data/train-labels-idx1-ubyte" & cBatchN 32 & cLabelToNDArray) -testingData :: MonadResource m => Stream (Of (ArrayF, ArrayF)) m ()+testingData :: MonadResource m => Stream (Of (ArrayF, ArrayF)) m Int testingData = S.zip     (sourceImages "examples/data/t10k-images-idx3-ubyte" & cBatchN 1 & cImageToNDArray      )-    (sourceLabels "examples/data/t10k-labels-idx1-ubyte" & cBatchN 1 & cLabelToOnehotNDArray)+    (sourceLabels "examples/data/t10k-labels-idx1-ubyte" & cBatchN 1 & cLabelToNDArray)  newtype Batched a = Batched { _batch :: NV.Vector a } 
examples/mnist/Parse.hs view
@@ -38,19 +38,19 @@ label :: AP.Parser Label label = fromIntegral <$> AP.anyWord8 -sourceImages :: MonadResource m => FilePath -> Stream (Of Image) m ()+sourceImages :: MonadResource m => FilePath -> Stream (Of Image) m Int sourceImages fp = do   (result, rest)<- lift $ APS.parse header (BSS.readFile fp)   case result of-    Left (HeaderImg _ w h) -> void $ APS.parsed (image w h) rest-    _ -> throwM NotImageFile+    Left (HeaderImg n w h) -> APS.parsed (image w h) rest >> return n+    _ -> effect $ throwM NotImageFile -sourceLabels :: MonadResource m => FilePath -> Stream (Of Label) m ()+sourceLabels :: MonadResource m => FilePath -> Stream (Of Label) m Int sourceLabels fp = do   (result, rest)<- lift $ APS.parse header (BSS.readFile fp)   case result of-    Left (HeaderLbl _) -> void $ APS.parsed label rest-    _ -> throwM NotImageFile+    Left (HeaderLbl n) -> APS.parsed label rest >> return n+    _ -> effect $ throwM NotImageFile  data Exc = NotImageFile | NotLabelFile     deriving (Show, Typeable)
+ examples/mnist/lenet.hs view
@@ -0,0 +1,123 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE FlexibleContexts #-}+module Main where++import MXNet.Core.Base hiding (variable, convolution, fullyConnected)+import qualified MXNet.Core.Base.NDArray as A+import qualified MXNet.Core.Base.Internal.TH.NDArray as A+import qualified MXNet.Core.Base.Symbol as S+import qualified MXNet.Core.Base.Internal.TH.Symbol as S+import qualified Data.HashMap.Strict as M+import Control.Monad (forM_, void)+import qualified Streaming.Prelude as SR+import qualified Data.Vector.Storable as SV+import Control.Monad.IO.Class+import Control.Monad.Trans.Resource+import System.IO (hFlush, stdout)+import MXNet.NN+import MXNet.NN.Utils+import MXNet.NN.Layer+import Dataset++-- # first conv+-- conv1 = mx.symbol.Convolution(data=data, kernel=(5,5), num_filter=20)+-- tanh1 = mx.symbol.Activation(data=conv1, act_type="tanh")+-- pool1 = mx.symbol.Pooling(data=tanh1, pool_type="max", kernel=(2,2), stride=(2,2))+-- # second conv+-- conv2 = mx.symbol.Convolution(data=pool1, kernel=(5,5), num_filter=50)+-- tanh2 = mx.symbol.Activation(data=conv2, act_type="tanh")+-- pool2 = mx.symbol.Pooling(data=tanh2, pool_type="max", kernel=(2,2), stride=(2,2))+-- # first fullc+-- flatten = mx.symbol.Flatten(data=pool2)+-- fc1 = mx.symbol.FullyConnected(data=flatten, num_hidden=500)+-- tanh3 = mx.symbol.Activation(data=fc1, act_type="tanh")+-- # second fullc+-- fc2 = mx.symbol.FullyConnected(data=tanh3, num_hidden=num_classes)+-- # loss+-- lenet = mx.symbol.SoftmaxOutput(data=fc2, name='softmax')++neural :: IO SymbolF+neural = do+    x  <- variable "x"+    y  <- variable "y"++    v1 <- convolution "conv1" x [5,5] 20 nil+    a1 <- S.activation "conv1-a" v1 "tanh"+    p1 <- S.pooling "conv1-p" a1 "(2,2)" "max" nil++    v2 <- convolution "conv2" p1 [5,5] 50 nil+    a2 <- S.activation "conv2-a" v2 "tanh"+    p2 <- S.pooling "conv2-p" a2 "(2,2)" "max" nil++    fl <- S.flatten "flatten" p2++    v3 <- fullyConnected "fc1" fl 500 nil+    a3 <- S.activation "fc1-a" v3 "tanh"++    v4 <- fullyConnected "fc2" a3 10 nil+    a4 <- S.softmaxoutput "softmax" v4 y nil+    return $ S.Symbol a4++range :: Int -> [Int]+range = enumFromTo 1++default_initializer :: DType a => Initializer a+default_initializer cxt shape = A.NDArray <$> A.random_normal +                                        (add @"loc" 0 $ +                                         add @"scale" 1 $ +                                         add @"shape" (formatShape shape) $ +                                         add @"ctx" (formatContext cxt) nil)+    +optimizer :: DType a => Optimizer a+optimizer _ v g = A.NDArray <$> A.sgd_update (A.getHandle v) (A.getHandle g) 0.0002 nil++main :: IO ()+main = do+    -- call mxListAllOpNames can ensure the MXNet itself is properly initialized+    -- i.e. MXNet operators are registered in the NNVM+    _    <- mxListAllOpNames+    net  <- neural+    sess <- initialize net $ Config { +                _cfg_placeholders = M.singleton "x" [1,1,28,28],+                _cfg_initializers = M.empty,+                _cfg_default_initializer = default_initializer,+                _cfg_context = contextGPU+            }++    runResourceT $ train sess $ do +        liftIO $ putStrLn $ "[Train] "+        let index = SR.enumFrom (1 :: Int)+        forM_ (range 50) $ \ind -> do+            liftIO $ putStrLn $ "iteration " ++ show ind+            total <- SR.effects trainingData+            _ <- flip SR.mapM_ (SR.zip index trainingData) $ \(i, (x, y)) -> do+                 liftIO $ do+                    putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show total+                    hFlush stdout+                 fit optimizer net $ M.fromList [("x", x), ("y", y)]+            liftIO $ putStr "\r\ESC[K"+        +        liftIO $ putStrLn $ "[Test] "+        total <- SR.effects testingData+        result<- SR.toList_ $ void $ flip SR.mapM (SR.zip index testingData) $ \(i, (x, y)) -> do +            liftIO $ do +                putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show total+                hFlush stdout+            [y'] <- forwardOnly net (M.fromList [("x", Just x), ("y", Nothing)])+            ind1 <- liftIO $ items y+            ind2 <- liftIO $ argmax y' >>= items+            return (ind1, ind2)+        liftIO $ putStr "\r\ESC[K"++        let (ls,ps) = unzip result+            ls_unbatched = mconcat ls+            ps_unbatched = mconcat ps+            total_test_items = SV.length ls_unbatched+            correct = SV.length $ SV.filter id $ SV.zipWith (==) ls_unbatched ps_unbatched+        liftIO $ putStrLn $ "Accuracy: " ++ show correct ++ "/" ++ show total_test_items+  +  where+    argmax :: ArrayF -> IO ArrayF+    argmax ys = A.NDArray <$> A.argmax (A.getHandle ys) (add @"axis" 1 nil)
examples/mnist/mnist.hs view
@@ -2,62 +2,64 @@ {-# LANGUAGE TypeApplications #-} module Main where -import MXNet.Core.Base+import MXNet.Core.Base hiding (variable, convolution, fullyConnected) import qualified MXNet.Core.Base.NDArray as A import qualified MXNet.Core.Base.Internal.TH.NDArray as A+import qualified MXNet.Core.Base.Symbol as S+import qualified MXNet.Core.Base.Internal.TH.Symbol as S import qualified Data.HashMap.Strict as M-import Control.Monad (forM_)+import Control.Monad (forM_, void) import qualified Streaming.Prelude as SR import qualified Data.Vector.Storable as SV-import Data.List (intersperse) import Control.Monad.IO.Class import Control.Monad.Trans.Resource import MXNet.NN+import MXNet.NN.Utils+import MXNet.NN.Layer import Dataset  neural :: IO SymbolF neural = do-    x  <- variable "x"  :: IO SymbolF -    y  <- variable "y"  :: IO SymbolF-    w1 <- variable "w1" :: IO SymbolF-    b1 <- variable "b1" :: IO SymbolF-    v1 <- fullyConnected x w1 b1 128-    a1 <- activation v1 "relu"-    w2 <- variable "w2" :: IO SymbolF-    b2 <- variable "b2" :: IO SymbolF-    v2 <- fullyConnected a1 w2 b2 10-    a2 <- softmaxOutput v2 y -    return a2+    x  <- variable "x"+    y  <- variable "y"+    v1 <- fullyConnected "fc1" x 128 nil+    a1 <- S.activation "fc1-a" v1 "relu"+    v2 <- fullyConnected "fc2" a1 10 nil+    a2 <- S.softmaxoutput "softmax" v2 y nil+    return $ S.Symbol a2  range :: Int -> [Int] range = enumFromTo 1 -default_initializer :: DType a => [Int] -> IO (NDArray a)-default_initializer shape = A.NDArray <$> A.random_normal (add @"loc" 0 $ add @"scale" 1 $ add @"shape" formatedShape nil)-  where-    formatedShape = concat $ ["("] ++ intersperse "," (map show shape) ++ [")"]-    -optimizer :: DType a => NDArray a -> NDArray a -> IO (NDArray a)-optimizer v g = A.NDArray <$> (A.sgd_update (A.getHandle v) (A.getHandle g) 0.01 nil)+default_initializer :: DType a => Initializer a+default_initializer cxt shape = A.NDArray <$> A.random_normal +                                        (add @"loc" 0 $ +                                         add @"scale" 1 $ +                                         add @"shape" (formatShape shape) $ +                                         add @"ctx" (formatContext cxt) nil)    +optimizer :: DType a => Optimizer a+optimizer _ v g = A.NDArray <$> (A.sgd_update (A.getHandle v) (A.getHandle g) 0.01 nil)  main :: IO () main = do     -- call mxListAllOpNames can ensure the MXNet itself is properly initialized     -- i.e. MXNet operators are registered in the NNVM-    _  <- mxListAllOpNames-    net <- neural-    params <- initialize net $ Config { +    _    <- mxListAllOpNames+    net  <- neural+    sess <- initialize net $ Config {                  _cfg_placeholders = M.singleton "x" [32,28,28],                 _cfg_initializers = M.empty,-                _cfg_default_initializer = default_initializer-              }-    result <- runResourceT $ train params contextCPU $ do +                _cfg_default_initializer = default_initializer,+                _cfg_context = contextGPU+            }+    result <- runResourceT $ train sess $ do          liftIO $ putStrLn $ "[Train] "         forM_ (range 5) $ \ind -> do             liftIO $ putStrLn $ "iteration " ++ show ind             SR.mapM_ (\(x, y) -> fit optimizer net $ M.fromList [("x", x), ("y", y)]) trainingData+         liftIO $ putStrLn $ "[Test] "-        SR.toList_ $ flip SR.mapM testingData $ \(x, y) -> do +        SR.toList_ $ void $ flip SR.mapM testingData $ \(x, y) -> do              [y'] <- forwardOnly net (M.fromList [("x", Just x), ("y", Nothing)])             ind1 <- liftIO $ argmax y  >>= items             ind2 <- liftIO $ argmax y' >>= items
mxnet-nn.cabal view
@@ -1,5 +1,5 @@ name:                       mxnet-nn-version:                    0.0.1.1+version:                    0.0.1.2 synopsis:                   Train a neural network with MXNet in Haskell. description:                High level APIs to rain a neural network with MXNet in Haskell. homepage:                   http://github.com/pierric/mxnet-nn@@ -14,6 +14,8 @@  Library     exposed-modules:        MXNet.NN+                            MXNet.NN.Utils+                            MXNet.NN.Layer     other-modules:     hs-source-dirs:         src     ghc-options:            -Wall@@ -28,6 +30,29 @@  Executable mnist     main-is:                mnist.hs+    other-modules:          Parse Dataset+    hs-source-dirs:         examples/mnist+    ghc-options:            -Wall+    default-language:       Haskell2010+    build-depends:          base >= 4.7 && < 5.0+                          , mxnet >= 0.2.0.0+                          , unordered-containers >= 0.2.8+                          , attoparsec >= 0.13+                          , attoparsec-binary >= 0.2+                          , vector >= 0.12+                          , bytestring >= 0.10+                          , resourcet >= 1.1.8+                          , exceptions >= 0.8.3+                          , mmorph >= 1.0.9+                          , mtl >= 2.2.0+                          , streaming >= 0.1.4.5+                          , streaming-utils >= 0.1.4.5+                          , streaming-bytestring >= 0.1.4.5+                          , ghc-prim+                          , mxnet-nn++Executable lenet+    main-is:                lenet.hs     other-modules:          Parse Dataset     hs-source-dirs:         examples/mnist     ghc-options:            -Wall
src/MXNet/NN.hs view
@@ -1,9 +1,13 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE TemplateHaskell #-} module MXNet.NN (     Parameter(..),     Config(..),+    Session(..),     Exc(..),     Initializer,     Optimizer,@@ -12,42 +16,50 @@     inferShape,     initialize,     fit,-    forwardOnly+    forwardOnly,+    getContext,+    sess_param,+    sess_context, ) where -import MXNet.Core.Base hiding (bind, context)+import MXNet.Core.Base hiding (bind, context, (^.)) import MXNet.Core.Base.Internal import qualified MXNet.Core.Base.NDArray as A import qualified MXNet.Core.Base.Symbol as S import qualified MXNet.Core.Base.Executor as E import qualified MXNet.Core.Types.Internal as MXI+import qualified MXNet.Core.Base.Internal.TH.NDArray as MXI import qualified Data.HashMap.Strict as M import Data.Typeable-import qualified Control.Monad.State as ST-import Data.Maybe (isJust, fromJust)+import qualified Control.Monad.State.Strict as ST+import Data.Maybe (isJust, fromJust, maybe) import Control.Monad (when) import Control.Monad.IO.Class (MonadIO, liftIO) import Control.Monad.Trans.Resource (MonadThrow(..)) import Control.Exception.Base (Exception)-import Control.Lens (traverseOf, _1)+import Control.Lens (makeLenses, traverseOf, use)  -- | A parameter is two 'NDArray' to back a 'Symbol' data Parameter a = Parameter { _param_in :: NDArray a, _param_grad :: NDArray a }     deriving Show --- | TrainM is a 'StateT' monad, where the state is all the 'Parameters' and a 'Context'-type TrainM a m = ST.StateT (M.HashMap String (Parameter a), Context) m+-- | Session is all the 'Parameters' and a 'Context'+-- type Session a = (M.HashMap String (Parameter a), Context)+data Session a = Session { _sess_param :: !(M.HashMap String (Parameter a)), _sess_context :: !Context }+makeLenses ''Session+-- | TrainM is a 'StateT' monad+type TrainM a m = ST.StateT (Session a) m  -- | Initializer is about how to create a NDArray from a given shape.  --  -- Usually, it can be a wrapper of MXNet operators, such as @random_uniform@, @random_normal@,  -- @random_gamma@, etc..-type Initializer a = [Int] -> IO (NDArray a)-type Optimizer a = NDArray a -> NDArray a -> IO (NDArray a)+type Initializer a = Context -> [Int] -> IO (NDArray a)+type Optimizer a = Context -> NDArray a -> NDArray a -> IO (NDArray a)      -- | Execute the 'TrainM' monad-train :: (DType a, Monad m) => M.HashMap String (Parameter a) -> Context -> TrainM a m r -> m r-train param context = flip ST.evalStateT (param, context)+train :: (DType a, Monad m) => Session a -> TrainM a m r -> m r+train = flip ST.evalStateT  -- | infer the shapes of all the symbols in a symbolic neural network inferShape :: DType a => Symbol a -> M.HashMap String (NDArray a) -> IO (M.HashMap String [Int])@@ -73,67 +85,88 @@ data Config a = Config {     _cfg_placeholders :: M.HashMap String [Int],     _cfg_initializers :: M.HashMap String (Initializer a),-    _cfg_default_initializer :: Initializer a+    _cfg_default_initializer :: Initializer a,+    _cfg_context :: Context }  -- | initialize all parameters-initialize :: DType a => Symbol a -> Config a -> IO (M.HashMap String (Parameter a))+initialize :: DType a => Symbol a -> Config a -> IO (Session a) initialize sym config = do     let spec1 = M.difference (_cfg_placeholders config) (_cfg_initializers config)         spec2 = _cfg_initializers config         dinit = _cfg_default_initializer config-    placeholder  <- mapM zeros spec1+        cxt   = _cfg_context config+    placeholder  <- mapM (\shp -> makeEmptyNDArray shp cxt False) spec1     inp_with_shp <- inferShape sym placeholder-    M.traverseWithKey (init_with_random_normal placeholder spec2 dinit) inp_with_shp+    args <- M.traverseWithKey (init_with_random_normal placeholder spec2 dinit) inp_with_shp+    return $ Session args cxt   where     init_with_random_normal placeholder spec2 dinit inp shp = do         case M.lookup inp placeholder of-            Just in_arg -> return $ Parameter in_arg (A.NDArray MXI.nullNDArrayHandle)+            Just in_arg -> do+                nullarg <- MXI.nullNDArrayHandle+                return $ Parameter in_arg (A.NDArray nullarg)             Nothing -> do                 arg_in <- case M.lookup inp spec2 of-                    Just cinit -> cinit shp-                    Nothing    -> dinit shp-                arg_gr <- zeros shp+                    Just cinit -> cinit (_cfg_context config) shp+                    Nothing    -> dinit (_cfg_context config) shp+                arg_gr <- makeEmptyNDArray shp (_cfg_context config) False                 return $ Parameter arg_in arg_gr  -- | bind the symbolic network with actual parameters-bind :: DType a => Symbol a -> M.HashMap String (Parameter a) -> Context -> Bool -> IO (Executor a)-bind net args Context{..} train_ = do-    names <- listInputs net-    exec_handle <- checked $ mxExecutorBind (S.getHandle net) deviceType deviceId-        (fromIntegral (M.size args))+bind :: (DType a, MonadIO m) => Symbol a -> Bool -> TrainM a m (Executor a)+bind net train_ = do+    args <- use sess_param+    Context{..} <- use sess_context+    exec_handle <- liftIO $ do+        names <- listInputs net+        nullarg <- MXI.nullNDArrayHandle         -- the parameters to bind should be arranged in the same order as the names-        (map (A.getHandle . _param_in) $ map (args M.!) names)-        (if train_-            then map (A.getHandle . _param_grad) $ map (args M.!) names-            else replicate (M.size args) MXI.nullNDArrayHandle)-        (replicate (M.size args) 1)-        0 []+        let arg_num = fromIntegral (M.size args)+            arg_in  = map (A.getHandle . _param_in) $ map (args M.!) names+            arg_gr  = if train_ +                        then map (A.getHandle . _param_grad) $ map (args M.!) names+                        else replicate (M.size args) nullarg+            arg_gr_req = replicate (M.size args) 1 -    makeExecutor exec_handle+        checked $ mxExecutorBind (S.getHandle net) deviceType deviceId+                                            arg_num arg_in arg_gr arg_gr_req +                                            0 []+    return $ E.Executor exec_handle  -- | single step train. Must provide all the placeholders. fit :: (DType a, MonadIO m, MonadThrow m) => Optimizer a -> Symbol a -> M.HashMap String (NDArray a) -> TrainM a m () fit opt net datAndLbl = do     shps <- liftIO $ inferShape net datAndLbl-    modifyT . traverseOf _1 $ M.traverseWithKey $ \k p -> do+    modifyT . traverseOf sess_param $ M.traverseWithKey $ \k p -> do         let ishp = shps M.! k         case M.lookup k datAndLbl of-            Just a  -> return $ p {_param_in = a}+            Just a  -> liftIO $ update_param (Left a) p             Nothing -> do                 (_, pshp1) <- liftIO $ ndshape (_param_in p)                 (_, pshp2) <- liftIO $ ndshape (_param_grad p)                 when (ishp /= pshp1 || ishp /= pshp2) (throwM $ MismatchedShape k)                 return p-    (params, context) <- ST.get-    liftIO $ do-        exec <- bind net params context True+    exec <- bind net True+    liftIO $ do          checked $ mxExecutorForward (E.getHandle exec) 1-        backward exec-    modifyT . traverseOf _1  $ M.traverseWithKey $ \ k v -> do+        checked $ mxExecutorBackward (E.getHandle exec) 0 []+        -- forward/backward are asynchronised operation in mxnet, in a+        -- sense that only opcodes are pushed onto an internal execution +        -- stack, and there is a executor running in a separate thread.+        -- It is possible that an OOM of CPU memory occurs, if 'fit' are +        -- called so fast that too many opcodes and data on the stack, +        -- as described in issue #1+        checked $ mxNDArrayWaitAll+    cxt <- use sess_context+    modifyT . traverseOf sess_param  $ M.traverseWithKey $ \ k v -> do         if (not $ M.member k datAndLbl)-            then do new_in <- liftIO $ opt (_param_in v) (_param_grad v) -                    return $ v {_param_in = new_in}+            then do new_in <- liftIO $ opt cxt (_param_in v) (_param_grad v) +                    -- must evaluate the new parameter to WHNF+                    -- otherwise, the old _param_in is retained.+                    -- if context is GPU, then OOM will soon +                    -- occur, as described in issue #2+                    return $! v {_param_in = new_in}             else return v  -- | forward only. Must provide all the placeholders, setting the data to @Just xx@, and set label to @Nothing@.@@ -142,23 +175,47 @@ forwardOnly :: (DType a, MonadIO m, MonadThrow m) => Symbol a -> M.HashMap String (Maybe (NDArray a)) -> TrainM a m [NDArray a] forwardOnly net dat = do     shps <- liftIO $ inferShape net (M.map fromJust $ M.filter isJust dat)-    modifyT . traverseOf _1 $ M.traverseWithKey $ \k p -> do+    modifyT . traverseOf sess_param $ M.traverseWithKey $ \k p -> do         let ishp = shps M.! k         case M.lookup k dat of-            Just (Just a) ->-                return $ p {_param_in = a}-            Just Nothing  -> do-                dummy <- liftIO $ zeros ishp-                return $ p {_param_in = dummy}+            Just a -> liftIO $ update_param (maybe (Right ishp) Left a) p              Nothing -> do                 (_, pshp) <- liftIO $ ndshape (_param_in p)                 when (ishp /= pshp) (throwM $ MismatchedShape k)                 return p-    (params, context) <- ST.get+    exec <- bind net False     liftIO $ do-        exec <- bind net params context False         checked $ mxExecutorForward (E.getHandle exec) 0+        -- for the same reason in 'fit'.+        checked $ mxNDArrayWaitAll         getOutputs exec++update_param :: DType a => Either (NDArray a) [Int] -> Parameter a -> IO (Parameter a)+update_param (Left a) p = do+    src_cxt <- A.context a+    src_shp <- snd <$> A.ndshape a+    dst_cxt <- A.context (_param_in p)+    dst_shp <- snd <$> A.ndshape (_param_in p)+    case (src_cxt == dst_cxt, src_shp == dst_shp) of+        (True , True) -> return $ p {_param_in = a}+        (False, True) -> do+            MXI._copyto' (A.getHandle a) [A.getHandle (_param_in p)] :: IO ()+            return p+        _ -> do+            a_copy <- makeEmptyNDArray src_shp dst_cxt False+            MXI._copyto' (A.getHandle a) [A.getHandle a_copy] :: IO ()+            return $! p {_param_in = a_copy}    +update_param (Right src_shp) p = do+    dst_cxt <- A.context (_param_in p)+    dst_shp <- snd <$> A.ndshape (_param_in p)+    if src_shp == dst_shp +        then return p+        else do+            dummy <- makeEmptyNDArray src_shp dst_cxt False+            return $! p {_param_in = dummy}++getContext :: Monad m => TrainM a m Context+getContext = use sess_context  -- | Possible exception in 'TrainM' data Exc = MismatchedShape String
+ src/MXNet/NN/Layer.hs view
@@ -0,0 +1,32 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE FlexibleContexts #-}++module MXNet.NN.Layer where++import MXNet.Core.Types.Internal+import MXNet.Core.Base.HMap+import qualified MXNet.Core.Base.Internal.TH.Symbol as S+import qualified MXNet.Core.Base.Internal as I+import MXNet.NN.Utils++variable :: String -> IO SymbolHandle+variable = I.checked . I.mxSymbolCreateVariable++convolution :: (MatchKVList kvs '["stride" ':= String, "dilate" ':= String, "pad" ':= String,+                                  "num_group" ':= Int, "workspace" ':= Int, "no_bias" ':= Bool,+                                  "cudnn_tune" ':= String, "cudnn_off" ':= Bool, "layout" ':= String],+                ShowKV kvs)+            => String -> SymbolHandle -> [Int] -> Int -> HMap kvs -> IO SymbolHandle+convolution name dat kernel_shape num_filter args = do+    w <- variable (name ++ "-w")+    b <- variable (name ++ "-b")+    S.convolution name dat w b (formatShape kernel_shape) num_filter args++fullyConnected :: (MatchKVList kvs '["no_bias" ':= Bool, "flatten" ':= Bool], ShowKV kvs) +               => String -> SymbolHandle -> Int -> HMap kvs -> IO SymbolHandle+fullyConnected name dat num_neuron args = do+    w <- variable (name ++ "-w")+    b <- variable (name ++ "-b")+    S.fullyconnected name dat w b num_neuron args
+ src/MXNet/NN/Utils.hs view
@@ -0,0 +1,16 @@+{-# LANGUAGE RecordWildCards #-}+module MXNet.NN.Utils where++import MXNet.Core.Base.DType+import Data.List (intersperse)++formatShape :: [Int] -> String+formatShape shape = concat $ ["("] ++ intersperse "," (map show shape) ++ [")"]++formatContext :: Context -> String+formatContext Context{..} = getDeviceName deviceType ++ "(" ++ show deviceId ++ ")"+  where +    getDeviceName 1 = "cpu"+    getDeviceName 2 = "gpu"+    getDeviceName 3 = "cpu_pinned"+    getDeviceName _ = error "formatContext: unknown device type"