diff --git a/examples/mnist/Dataset.hs b/examples/mnist/Dataset.hs --- a/examples/mnist/Dataset.hs +++ b/examples/mnist/Dataset.hs @@ -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 } diff --git a/examples/mnist/Parse.hs b/examples/mnist/Parse.hs --- a/examples/mnist/Parse.hs +++ b/examples/mnist/Parse.hs @@ -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) diff --git a/examples/mnist/lenet.hs b/examples/mnist/lenet.hs new file mode 100644 --- /dev/null +++ b/examples/mnist/lenet.hs @@ -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) diff --git a/examples/mnist/mnist.hs b/examples/mnist/mnist.hs --- a/examples/mnist/mnist.hs +++ b/examples/mnist/mnist.hs @@ -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 diff --git a/mxnet-nn.cabal b/mxnet-nn.cabal --- a/mxnet-nn.cabal +++ b/mxnet-nn.cabal @@ -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 diff --git a/src/MXNet/NN.hs b/src/MXNet/NN.hs --- a/src/MXNet/NN.hs +++ b/src/MXNet/NN.hs @@ -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 diff --git a/src/MXNet/NN/Layer.hs b/src/MXNet/NN/Layer.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/Layer.hs @@ -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 diff --git a/src/MXNet/NN/Utils.hs b/src/MXNet/NN/Utils.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/Utils.hs @@ -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"