mxnet-nn 0.0.1.2 → 0.0.1.3
raw patch · 14 files changed
+519/−300 lines, 14 filesdep +transformers-basedep ~mtl
Dependencies added: transformers-base
Dependency ranges changed: mtl
Files
- examples/mnist/Dataset.hs +0/−66
- examples/mnist/DatasetVector.hs +55/−0
- examples/mnist/Parse.hs +0/−24
- examples/mnist/lenet.hs +33/−30
- examples/mnist/mnist.hs +0/−76
- mxnet-nn.cabal +9/−25
- src/MXNet/NN.hs +74/−74
- src/MXNet/NN/DataIter/Class.hs +31/−0
- src/MXNet/NN/DataIter/LazyVec.hs +84/−0
- src/MXNet/NN/EvalMetric.hs +88/−0
- src/MXNet/NN/Layer.hs +55/−5
- src/MXNet/NN/Optimizer.hs +64/−0
- src/MXNet/NN/Types.hs +24/−0
- src/MXNet/NN/Utils.hs +2/−0
− examples/mnist/Dataset.hs
@@ -1,66 +0,0 @@-{-# LANGUAGE KindSignatures #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE OverloadedLists #-}-module Dataset where--import MXNet.Core.Base-import qualified MXNet.Core.Base.NDArray as A-import qualified MXNet.Core.Base.Internal.TH.NDArray as MXI-import Data.Function ((&))-import Streaming-import Streaming.Prelude (Of(..))-import qualified Streaming.Prelude as S-import Control.Monad.Trans.Resource (MonadResource(..))-import qualified Data.Vector as NV-import qualified Data.Vector.Storable as SV--import Parse--type SymbolF = Symbol Float-type ArrayF = NDArray Float--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)-mappedOf = S.mapM--cImageToNDArray :: MonadIO m => StreamProc (Batched Image) ArrayF m-cImageToNDArray = mappedOf $ \dat -> liftIO $ do- let sz = size 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 <- 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 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 Int-trainingData = S.zip- (sourceImages "examples/data/train-images-idx3-ubyte" & cBatchN 32 & cImageToNDArray )- (sourceLabels "examples/data/train-labels-idx1-ubyte" & cBatchN 32 & cLabelToNDArray)--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 & cLabelToNDArray)--newtype Batched a = Batched { _batch :: NV.Vector a }--size :: Batched a -> Int-size (Batched b) = NV.length b-
+ examples/mnist/DatasetVector.hs view
@@ -0,0 +1,55 @@+module DatasetVector where + +import MXNet.Core.Base +import Data.Typeable +import Control.Monad.Trans.Resource (MonadResource(..), MonadThrow(..)) +import Control.Monad.IO.Class (liftIO) +import Control.Exception.Base +import qualified Data.Vector as V +import qualified Data.Vector.Storable as VS +import qualified Data.ByteString as BS +import Data.Attoparsec.ByteString as AP + +import MXNet.NN.DataIter.LazyVec as VL +import Parse + +type SymbolF = Symbol Float +type ArrayF = NDArray Float + +loadTrainingData :: MonadResource m => m (VL.LVec (ArrayF, ArrayF)) +loadTrainingData = do + v1 <- sourceImages "examples/data/train-images-idx3-ubyte" >>= liftIO . batch 128 + v2 <- sourceLabels "examples/data/train-labels-idx1-ubyte" >>= liftIO . batch 128 + return $ VL.zip (VL.map cImageToNDArray v1) (VL.map cLabelToNDArray v2) + +loadTestingData :: MonadResource m => m (VL.LVec (ArrayF, ArrayF)) +loadTestingData = do + v1 <- sourceImages "examples/data/t10k-images-idx3-ubyte" >>= liftIO . batch 1 + v2 <- sourceLabels "examples/data/t10k-labels-idx1-ubyte" >>= liftIO . batch 1 + return $ VL.zip (VL.map cImageToNDArray v1) (VL.map cLabelToNDArray v2) + +sourceImages :: MonadResource m => FilePath -> m (VL.LVec Image) +sourceImages = parseFile $ do + HeaderImg n w h <- header + count n (image w h) + +sourceLabels :: MonadResource m => FilePath -> m (VL.LVec Label) +sourceLabels = parseFile $ do + HeaderLbl n <- header + count n label + +parseFile :: MonadResource m => Parser [a] -> FilePath -> m (VL.LVec a) +parseFile parser fp = do + content <- liftIO $ BS.readFile fp + case AP.parseOnly parser content of + Left msg -> throwM $ ParseError msg + Right rt -> return $ VL.fromVec $ V.fromList rt + +cImageToNDArray :: V.Vector Image -> IO ArrayF +cImageToNDArray dat = makeNDArray [V.length dat, 1, 28, 28] contextCPU (VS.concat $ V.toList dat) +cLabelToNDArray :: V.Vector Label -> IO ArrayF +cLabelToNDArray dat = makeNDArray [V.length dat] contextCPU (V.convert $ V.map fromIntegral dat) + +data Exc = ParseError String + deriving (Show, Typeable) +instance Exception Exc
examples/mnist/Parse.hs view
@@ -1,15 +1,9 @@ module Parse where -import Streaming import Data.Attoparsec.ByteString as AP import Data.Attoparsec.Binary as AP-import Data.Attoparsec.ByteString.Streaming as APS-import qualified Data.ByteString.Streaming as BSS import qualified Data.ByteString.Internal as BS import qualified Data.Vector.Storable as SV-import Control.Exception.Base-import Control.Monad.Trans.Resource (MonadResource(..), MonadThrow(..))-import Data.Typeable type Image = SV.Vector Float type Label = Int@@ -37,21 +31,3 @@ label :: AP.Parser Label label = fromIntegral <$> AP.anyWord8--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 n w h) -> APS.parsed (image w h) rest >> return n- _ -> effect $ throwM NotImageFile--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 n) -> APS.parsed label rest >> return n- _ -> effect $ throwM NotImageFile--data Exc = NotImageFile | NotLabelFile- deriving (Show, Typeable)-instance Exception Exc
examples/mnist/lenet.hs view
@@ -4,14 +4,13 @@ {-# LANGUAGE FlexibleContexts #-} module Main where -import MXNet.Core.Base hiding (variable, convolution, fullyConnected)+import MXNet.Core.Base (DType, contextCPU, contextGPU, mxListAllOpNames)+import MXNet.Core.Base.HMap 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@@ -19,8 +18,11 @@ import MXNet.NN import MXNet.NN.Utils import MXNet.NN.Layer-import Dataset+import MXNet.NN.EvalMetric+import MXNet.NN.DataIter.Class +import DatasetVector+ -- # first conv -- conv1 = mx.symbol.Convolution(data=data, kernel=(5,5), num_filter=20) -- tanh1 = mx.symbol.Activation(data=conv1, act_type="tanh")@@ -44,20 +46,20 @@ 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+ a1 <- activation "conv1-a" v1 Tanh+ p1 <- pooling "conv1-p" a1 [2,2] PoolingMax 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+ a2 <- activation "conv2-a" v2 Tanh+ p2 <- pooling "conv2-p" a2 [2,2] PoolingMax nil - fl <- S.flatten "flatten" p2+ fl <- flatten "flatten" p2 v3 <- fullyConnected "fc1" fl 500 nil- a3 <- S.activation "fc1-a" v3 "tanh"+ a3 <- activation "fc1-a" v3 Tanh v4 <- fullyConnected "fc2" a3 10 nil- a4 <- S.softmaxoutput "softmax" v4 y nil+ a4 <- softmaxoutput "softmax" v4 y nil return $ S.Symbol a4 range :: Int -> [Int]@@ -66,13 +68,10 @@ default_initializer :: DType a => Initializer a default_initializer cxt shape = A.NDArray <$> A.random_normal (add @"loc" 0 $ - add @"scale" 1 $ + add @"scale" 0.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@@ -83,31 +82,35 @@ _cfg_placeholders = M.singleton "x" [1,1,28,28], _cfg_initializers = M.empty, _cfg_default_initializer = default_initializer,- _cfg_context = contextGPU+ _cfg_context = contextCPU }+ optimizer <- makeOptimizer 0.002 nil :: IO (ADAM Float '[]) runResourceT $ train sess $ do ++ trainingData <- loadTrainingData+ testingData <- loadTestingData+ liftIO $ putStrLn $ "[Train] "- let index = SR.enumFrom (1 :: Int)- forM_ (range 50) $ \ind -> do+ forM_ (range 5) $ \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"+ metric <- newMetric CrossEntropy "CrossEntropy" ["y"]+ void $ forEachD_ni trainingData $ \((t,i), (x, y)) -> do+ liftIO $ do+ eval <- formatMetric metric+ putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show t ++ " " ++ eval+ hFlush stdout+ fitAndEval optimizer net (M.fromList [("x", x), ("y", y)]) metric+ liftIO $ putStrLn "" liftIO $ putStrLn $ "[Test] "- total <- SR.effects testingData- result<- SR.toList_ $ void $ flip SR.mapM (SR.zip index testingData) $ \(i, (x, y)) -> do + result <- forEachD_ni testingData $ \((t,i), (x, y)) -> do liftIO $ do - putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show total+ putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show t hFlush stdout [y'] <- forwardOnly net (M.fromList [("x", Just x), ("y", Nothing)])- ind1 <- liftIO $ items y- ind2 <- liftIO $ argmax y' >>= items+ ind1 <- liftIO $ A.items y+ ind2 <- liftIO $ argmax y' >>= A.items return (ind1, ind2) liftIO $ putStr "\r\ESC[K"
− examples/mnist/mnist.hs
@@ -1,76 +0,0 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE TypeApplications #-}-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 MXNet.NN-import MXNet.NN.Utils-import MXNet.NN.Layer-import Dataset--neural :: IO SymbolF-neural = do- 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 => 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- sess <- initialize net $ Config { - _cfg_placeholders = M.singleton "x" [32,28,28],- _cfg_initializers = M.empty,- _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_ $ 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- return (ind1, ind2)- let (ls,ps) = unzip result- ls_unbatched = mconcat ls- ps_unbatched = mconcat ps- total = SV.length ls_unbatched- correct = SV.length $ SV.filter id $ SV.zipWith (==) ls_unbatched ps_unbatched- putStrLn $ "Accuracy: " ++ show correct ++ "/" ++ show total- - where- argmax :: ArrayF -> IO ArrayF- argmax ys = A.NDArray <$> A.argmax (A.getHandle ys) (add @"axis" 1 nil)
mxnet-nn.cabal view
@@ -1,5 +1,5 @@ name: mxnet-nn-version: 0.0.1.2+version: 0.0.1.3 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,12 +14,18 @@ Library exposed-modules: MXNet.NN+ MXNet.NN.Types MXNet.NN.Utils MXNet.NN.Layer+ MXNet.NN.Optimizer+ MXNet.NN.EvalMetric+ MXNet.NN.DataIter.Class+ MXNet.NN.DataIter.LazyVec other-modules: hs-source-dirs: src ghc-options: -Wall default-language: Haskell2010+ default-extensions: GADTs, TypeFamilies build-depends: base >= 4.7 && < 5.0 , mxnet >= 0.2.0.0 , unordered-containers >= 0.2.8@@ -27,33 +33,11 @@ , vector >= 0.12 , mtl >= 2.2 , lens >= 4.12--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+ , transformers-base >= 0.4.4 Executable lenet main-is: lenet.hs- other-modules: Parse Dataset+ other-modules: Parse DatasetVector hs-source-dirs: examples/mnist ghc-options: -Wall default-language: Haskell2010
src/MXNet/NN.hs view
@@ -3,23 +3,22 @@ {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE BangPatterns #-}-{-# LANGUAGE TemplateHaskell #-} module MXNet.NN ( Parameter(..), Config(..), Session(..), Exc(..), Initializer,- Optimizer, TrainM, train, inferShape, initialize,- fit,+ fit, fitAndEval, forwardOnly, getContext, sess_param, sess_context,+ module MXNet.NN.Optimizer ) where import MXNet.Core.Base hiding (bind, context, (^.))@@ -33,30 +32,16 @@ import Data.Typeable import qualified Control.Monad.State.Strict as ST import Data.Maybe (isJust, fromJust, maybe)-import Control.Monad (when)+import Control.Monad (when, zipWithM_) import Control.Monad.IO.Class (MonadIO, liftIO) import Control.Monad.Trans.Resource (MonadThrow(..)) import Control.Exception.Base (Exception)-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+import Control.Lens (traverseOf, use, (^.)) --- | 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+import MXNet.NN.Types+import MXNet.NN.Optimizer+import MXNet.NN.EvalMetric --- | 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 = 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) => Session a -> TrainM a m r -> m r train = flip ST.evalStateT@@ -114,10 +99,24 @@ return $ Parameter arg_in arg_gr -- | bind the symbolic network with actual parameters-bind :: (DType a, MonadIO m) => Symbol a -> Bool -> TrainM a m (Executor a)-bind net train_ = do- args <- use sess_param+bind :: (DType a, MonadIO m, MonadThrow m) => Symbol a -> M.HashMap String (Maybe (NDArray a)) -> Bool -> TrainM a m (Executor a)+bind net dat train_ = do Context{..} <- use sess_context++ shps <- liftIO $ inferShape net (M.map fromJust $ M.filter isJust dat)+ modifyT . traverseOf sess_param $ M.traverseWithKey $ \k p -> do+ let ishp = shps M.! k+ case M.lookup k dat of+ Just a -> liftIO $ update_param (maybe (Right ishp) Left a) p+ Nothing -> do+ (_, pshp1) <- liftIO $ ndshape (_param_in p)+ when (ishp /= pshp1 ) (throwM $ MismatchedShape k)+ when train_ $ do+ (_, pshp2) <- liftIO $ ndshape (_param_grad p)+ when (ishp /= pshp2) (throwM $ MismatchedShape k)+ return p++ args <- use sess_param exec_handle <- liftIO $ do names <- listInputs net nullarg <- MXI.nullNDArrayHandle@@ -133,21 +132,36 @@ arg_num arg_in arg_gr arg_gr_req 0 [] return $ E.Executor exec_handle+ where+ 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} -- | 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 :: (DType a, MonadIO m, MonadThrow m, Optimizer opt, OptArgsCst opt g) + => opt a g -> Symbol a -> M.HashMap String (NDArray a) -> TrainM a m () fit opt net datAndLbl = do- shps <- liftIO $ inferShape net datAndLbl- modifyT . traverseOf sess_param $ M.traverseWithKey $ \k p -> do- let ishp = shps M.! k- case M.lookup k datAndLbl of- 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- exec <- bind net True+ exec <- bind net (M.map Just datAndLbl) True liftIO $ do checked $ mxExecutorForward (E.getHandle exec) 1 checked $ mxExecutorBackward (E.getHandle exec) 0 []@@ -158,10 +172,29 @@ -- called so fast that too many opcodes and data on the stack, -- as described in issue #1 checked $ mxNDArrayWaitAll- cxt <- use sess_context+ updateParameters opt datAndLbl++-- | single step train. Must provide all the placeholders.+-- After fitting, it also update the evaluation metric.+fitAndEval :: (DType a, MonadIO m, MonadThrow m, Optimizer opt, OptArgsCst opt g, EvalMetricMethod mth)+ => opt a g -> Symbol a -> M.HashMap String (NDArray a) -> Metric a mth -> TrainM a m ()+fitAndEval opt net datAndLbl metric = do+ exec <- bind net (M.map Just datAndLbl) True+ preds <- liftIO $ do + checked $ mxExecutorForward (E.getHandle exec) 1+ checked $ mxExecutorBackward (E.getHandle exec) 0 []+ checked $ mxNDArrayWaitAll+ getOutputs exec+ updateParameters opt datAndLbl+ let labels = map (datAndLbl M.!) (metric ^. metric_labelname)+ liftIO $ zipWithM_ (evaluate metric) preds labels++updateParameters :: (MonadIO m, Optimizer opt, OptArgsCst opt args) + => opt dtype args -> M.HashMap String any -> TrainM dtype m ()+updateParameters opt blacklist = do modifyT . traverseOf sess_param $ M.traverseWithKey $ \ k v -> do- if (not $ M.member k datAndLbl)- then do new_in <- liftIO $ opt cxt (_param_in v) (_param_grad v) + if (not $ M.member k blacklist)+ then do new_in <- liftIO $ optimize opt k (_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 @@ -174,45 +207,12 @@ -- Note that the batch size here can be different from that in the training phase. 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 sess_param $ M.traverseWithKey $ \k p -> do- let ishp = shps M.! k- case M.lookup k dat of- 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- exec <- bind net False+ exec <- bind net dat False liftIO $ do 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
+ src/MXNet/NN/DataIter/Class.hs view
@@ -0,0 +1,31 @@+{-# Language MultiParamTypeClasses, FunctionalDependencies #-} +{-# LANGUAGE FlexibleInstances #-} +module MXNet.NN.DataIter.Class where + +import GHC.Exts (Constraint) + +-- | Constraints on Dataset and the monad where the operation shall be ran. +type family DatasetConstraint (d :: * -> *) (m :: * -> *) :: Constraint + +-- | Abstract Dataset type class +class Dataset (d :: * -> *) where + -- | Create Dataset from `[]`. + -- note that depending on the instance, it may or may not work with infinitive list. + fromListD :: [e] -> d e + -- | Zip two Datasets + zipD :: d e1 -> d e2 -> d (e1, e2) + -- | Get number of elements + sizeD :: (DatasetConstraint d m, Monad m) => d e -> m Int + -- | Apply a function on each element of Dataset + forEachD :: (DatasetConstraint d m, Monad m) => d e -> (e -> m a) -> m [a] + + -- | Apply a function on each element of Dataset together with the element's index. + -- Note that the default implmentation assumes the Dataset can be created from a infinitive list. + forEachD_i :: (DatasetConstraint d m, Monad m) => d e -> ((Int, e) -> m a) -> m [a] + forEachD_i dat = forEachD (zipD (fromListD [1..]) dat) + + -- | Apply a function on each element of Dataset together with the total number of elements and the element's index. + forEachD_ni :: (DatasetConstraint d m, Monad m) => d e -> (((Int, Int), e) -> m a) -> m [a] + forEachD_ni dat proc = do + n <- sizeD dat + forEachD ((fromListD (replicate n n) `zipD` fromListD [1..n]) `zipD` dat) proc
+ src/MXNet/NN/DataIter/LazyVec.hs view
@@ -0,0 +1,84 @@+{-# Language MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleInstances #-} +module MXNet.NN.DataIter.LazyVec where + +import Prelude hiding (zip) +import Data.Vector (Vector) +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as VM +import Data.IORef +import Control.Monad +import Control.Monad.IO.Class +import Control.Exception.Base (assert) +import MXNet.NN.DataIter.Class + +data Lazy a = Direct a | Make (() -> IO a) + +instance Functor Lazy where + fmap f (Direct a) = Direct (f a) + fmap f (Make g) = Make (g >=> return . f) + +force :: Lazy a -> IO a +force (Direct a) = return a +force (Make f) = f () + +data LVec a = LVec { size :: Int, unLVec :: Lazy (Vector a)} + +fromVec :: Vector a -> LVec a +fromVec v = LVec (V.length v) (Direct v) + +toVec :: LVec a -> IO (Vector a) +toVec = force . unLVec + +batch :: Int -> LVec a -> IO (LVec (Vector a)) +batch chunksize vec = do + pos <- newIORef 0 + return $ case unLVec vec of + Direct v -> makeChunk pos v + Make f -> LVec new_vec_size . Make $ f >=> (toVec . makeChunk pos) + where + total = size vec + (quotient, remainder) = divMod total chunksize + new_vec_size = if remainder > 0 then quotient + 1 else quotient + makeChunk cur_pos vector = + assert (V.length vector == total) $ + LVec new_vec_size . Make $ \_ -> do + vec' <- VM.new new_vec_size + forM_ [0..new_vec_size-1] $ \ i -> do + j <- readIORef cur_pos + if j + chunksize >= total + then do + let rst = total - j + rnd = chunksize - rst + VM.write vec' i (V.slice j rst vector V.++ V.slice 0 rnd vector) + writeIORef cur_pos rnd + else do + VM.write vec' i (V.slice j chunksize vector) + writeIORef cur_pos (j+chunksize) + V.freeze vec' + +zip :: LVec a -> LVec b -> LVec (a,b) +zip (LVec n1 (Direct a)) (LVec n2 (Direct b)) = LVec (min n1 n2) (Direct (V.zip a b)) +zip (LVec n1 (Direct a)) (LVec n2 (Make f)) = LVec (min n1 n2) (Make (f >=> return . V.zip a)) +zip (LVec n1 (Make f)) (LVec n2 (Direct b)) = LVec (min n1 n2) (Make (f >=> return . flip V.zip b)) +zip (LVec n1 (Make f)) (LVec n2 (Make g)) = LVec (min n1 n2) (Make (\_ -> liftM2 V.zip (f ()) (g ()))) + +map :: (a -> IO b) -> LVec a -> LVec b +map f v = case fmap (V.mapM f) (unLVec v) of + Direct a -> LVec (size v) $ Make (\_ -> a) + Make m -> LVec (size v) $ Make (join . m) + +type instance DatasetConstraint LVec m = MonadIO m + +instance Dataset LVec where + fromListD = fromVec . V.fromList + zipD = zip + sizeD dat = return $ size dat + forEachD dat proc = do + vec <- liftIO $ toVec dat + ret <- V.mapM proc vec + return $ V.toList ret + + -- LVec does not support infinite stream, so we override the + -- default implementations + forEachD_i dat = forEachD (zipD (fromListD [1..size dat]) dat)
+ src/MXNet/NN/EvalMetric.hs view
@@ -0,0 +1,88 @@+{-# LANGUAGE TemplateHaskell #-} +module MXNet.NN.EvalMetric where + +import Data.IORef +import Control.Exception.Base (Exception) +import Control.Monad.Trans.Resource (MonadThrow(..)) +import Data.Typeable (Typeable) +import Control.Monad +import Control.Monad.IO.Class (MonadIO, liftIO) +import Text.Printf (printf) +import qualified Data.Vector.Storable as SV +import Control.Lens (makeLenses) +import MXNet.Core.Base +import qualified MXNet.Core.Base.NDArray as A +import qualified MXNet.Core.Base.Internal.TH.NDArray as A + +-- | Metric data +data Metric dytpe method = Metric { + _metric_name :: String, + _metric_labelname :: [String], + _metric_instance :: IORef Int, + _metric_sum :: IORef dytpe +} +makeLenses ''Metric + +-- | create a new metric data +newMetric :: (DType dtype, MonadIO m) => method -> String -> [String] -> m (Metric dtype method) +newMetric _ name labels = do + a <- liftIO $ newIORef 0 + b <- liftIO $ newIORef 0 + return $ Metric name labels a b + +-- | reset all information +resetMetric :: (DType dtype, MonadIO m) => Metric dtype method -> m () +resetMetric metric = liftIO $ do + writeIORef (_metric_sum metric) 0 + writeIORef (_metric_instance metric) 0 + +-- | get the metric +getMetric :: (DType dtype, MonadIO m) => Metric dtype method -> m Float +getMetric metric = do + s <- liftIO $ readIORef (_metric_sum metric) + n <- liftIO $ readIORef (_metric_instance metric) + return $ realToFrac s / fromIntegral n + +-- | format the metric as string +formatMetric :: (DType dtype, MonadIO m) => Metric dtype method -> m String +formatMetric metric = do + e <- getMetric metric + return $ printf "<%s: %0.3f>" (_metric_name metric) e + +-- | Abstract Evaluation type class +class EvalMetricMethod method where + evaluate :: DType dtype => Metric dtype method -> A.NDArray dtype -> A.NDArray dtype -> IO () + +-- | Basic evluation - cross entropy +data CrossEntropy = CrossEntropy +instance EvalMetricMethod CrossEntropy where + -- | evaluate the log-loss. + -- preds is of shape (batch_size, num_category), each element along the second dimension gives the probability of the category. + -- label is of shape (batch_size,), each element gives the category number. + evaluate metric preds label = do + (n1, shp1) <- A.ndshape preds + (n2, shp2) <- A.ndshape label + when (n1 /= 2 || n2 /= 1 || head shp1 /= head shp2) (throwM InvalidInput) + -- before call pick, we have to make sure preds and label + -- are in the same context + preds_may_copy <- do + c1 <- context preds + c2 <- context label + if c1 == c2 + then return preds + else do + (_, preds_shap) <- ndshape preds + preds_copy <- A.makeEmptyNDArray preds_shap c2 False + A._copyto' (A.getHandle preds) [A.getHandle preds_copy] :: IO () + return preds_copy + predprj <- A.pick (A.getHandle preds_may_copy) (A.getHandle label) nil + predlog <- A.log predprj + loss <- A.sum predlog nil >>= A.items . A.NDArray + modifyIORef (_metric_sum metric) (+ (negate $ loss SV.! 0)) + modifyIORef (_metric_instance metric) (+ head shp1) + +-- | Possible exceptions in evaluation. +data EvalMetricExc = InvalidInput + deriving (Show, Typeable) +instance Exception EvalMetricExc +
src/MXNet/NN/Layer.hs view
@@ -14,19 +14,69 @@ 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)+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) +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++data PoolingMethod = PoolingMax | PoolingAvg | PoolingSum++poolingMethodToStr :: PoolingMethod -> String+poolingMethodToStr PoolingMax = "max"+poolingMethodToStr PoolingAvg = "avg"+poolingMethodToStr PoolingSum = "sum"++pooling :: (MatchKVList kvs '["global_pool" ':= Bool,+ "cudnn_off" ':= Bool,+ "pooling_convention" ':= String,+ "stride" ':= String,+ "pad" ':= String]+ ,ShowKV kvs)+ => String -> SymbolHandle -> [Int] -> PoolingMethod -> HMap kvs -> IO SymbolHandle+pooling name input shape method args = S.pooling name input (formatShape shape) (poolingMethodToStr method) args++flatten :: String -> SymbolHandle -> IO SymbolHandle+flatten = S.flatten++data ActivationType = Relu | Sigmoid | Tanh | SoftRelu++activationTypeToStr :: ActivationType -> String+activationTypeToStr Relu = "relu"+activationTypeToStr Sigmoid = "sigmoid"+activationTypeToStr Tanh = "tanh"+activationTypeToStr SoftRelu = "softrelu"++activation :: String -> SymbolHandle -> ActivationType -> IO SymbolHandle+activation name input typ = S.activation name input (activationTypeToStr typ)++softmaxoutput :: (MatchKVList kvs '["grad_scale" ':= Float, + "ignore_label" ':= Float,+ "multi_output" ':= Bool, + "use_ignore" ':= Bool,+ "preserve_shape" ':= Bool, + "normalization" ':= String,+ "out_grad" ':= Bool, + "smooth_alpha" ':= Float],+ ShowKV kvs)+ => String -> SymbolHandle -> SymbolHandle -> HMap kvs -> IO SymbolHandle+softmaxoutput = S.softmaxoutput
+ src/MXNet/NN/Optimizer.hs view
@@ -0,0 +1,64 @@+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE ConstraintKinds #-}++module MXNet.NN.Optimizer (+ Optimizer(..),+ OptArgsCst,+ SGD, ADAM+) where++import qualified Data.HashMap.Strict as M+import MXNet.Core.Base.NDArray (NDArray)+import qualified MXNet.Core.Base.NDArray as A+import MXNet.Core.Base.HMap+import MXNet.Core.Base.Internal.TH.NDArray as A+import Data.IORef++-- | Constraint of using an optimizer+type OptArgsCst opt args = (ShowKV args, MatchKVList args (OptArgsList opt))++-- | Abstract Optimizer type class+class Optimizer (opt :: * -> [KV *] -> *) where+ -- | Specific constraints of the optimizer+ type OptArgsList opt :: [KV *]+ -- | make the optimizer+ makeOptimizer :: OptArgsCst opt args => Float -> HMap args -> IO (opt dtype args)+ -- | run the optimizer with the input & expected tensor+ optimize :: OptArgsCst opt args => opt dtype args -> String -> NDArray dytpe -> NDArray dtype -> IO (NDArray dtype)++-- | SGD optimizer+data SGD dtype args = SGD Float (HMap args)++instance Optimizer SGD where+ type OptArgsList SGD = '["wd" ':= Float,+ "rescale_grad" ':= Float,+ "clip_gradient" ':= Float]+ makeOptimizer lr args = return $ SGD lr args+ optimize (SGD lr args) _ weight gradient = A.NDArray <$> A.sgd_update (A.getHandle weight) (A.getHandle gradient) lr args++-- | ADAM optmizer+data ADAM dtype args = ADAM Float (HMap args) (IORef (M.HashMap String (NDArray dtype, NDArray dtype)))++instance Optimizer ADAM where+ type OptArgsList ADAM = '["beta1" ':= Float,+ "beta2" ':= Float,+ "epsilon" ':= Float,+ "wd" ':= Float,+ "rescale_grad" ':= Float,+ "clip_gradient" ':= Float]+ makeOptimizer lr args = do+ empty <- newIORef M.empty+ return $ ADAM lr args empty++ optimize (ADAM lr args emaref) symbol weight gradient = do+ ema <- readIORef emaref+ (moving_avg, moving_var) <- case M.lookup symbol ema of+ Nothing -> do+ avg <- A.zeros_like (A.getHandle weight) + var <- A.zeros_like (A.getHandle weight)+ writeIORef emaref (M.insert symbol (A.NDArray avg, A.NDArray var) ema)+ return (avg, var)+ Just (a,v) -> return (A.getHandle a, A.getHandle v)+ A.NDArray <$> adam_update (A.getHandle weight) (A.getHandle gradient) moving_avg moving_var lr args
+ src/MXNet/NN/Types.hs view
@@ -0,0 +1,24 @@+{-# LANGUAGE TemplateHaskell #-} +module MXNet.NN.Types where + +import Control.Lens (makeLenses) +import qualified Data.HashMap.Strict as M +import qualified Control.Monad.State.Strict as ST +import MXNet.Core.Base hiding (bind, context, (^.)) + +-- | A parameter is two 'NDArray' to back a 'Symbol' +data Parameter a = Parameter { _param_in :: NDArray a, _param_grad :: NDArray a } + deriving Show + +-- | 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 = Context -> [Int] -> IO (NDArray a)
src/MXNet/NN/Utils.hs view
@@ -4,9 +4,11 @@ import MXNet.Core.Base.DType import Data.List (intersperse) +-- | format a shape formatShape :: [Int] -> String formatShape shape = concat $ ["("] ++ intersperse "," (map show shape) ++ [")"] +-- | format a context formatContext :: Context -> String formatContext Context{..} = getDeviceName deviceType ++ "(" ++ show deviceId ++ ")" where