diff --git a/examples/mnist/Dataset.hs b/examples/mnist/Dataset.hs deleted file mode 100644 --- a/examples/mnist/Dataset.hs +++ /dev/null @@ -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 - diff --git a/examples/mnist/DatasetVector.hs b/examples/mnist/DatasetVector.hs new file mode 100644 --- /dev/null +++ b/examples/mnist/DatasetVector.hs @@ -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 diff --git a/examples/mnist/Parse.hs b/examples/mnist/Parse.hs --- a/examples/mnist/Parse.hs +++ b/examples/mnist/Parse.hs @@ -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 diff --git a/examples/mnist/lenet.hs b/examples/mnist/lenet.hs --- a/examples/mnist/lenet.hs +++ b/examples/mnist/lenet.hs @@ -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" diff --git a/examples/mnist/mnist.hs b/examples/mnist/mnist.hs deleted file mode 100644 --- a/examples/mnist/mnist.hs +++ /dev/null @@ -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) 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.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 diff --git a/src/MXNet/NN.hs b/src/MXNet/NN.hs --- a/src/MXNet/NN.hs +++ b/src/MXNet/NN.hs @@ -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 diff --git a/src/MXNet/NN/DataIter/Class.hs b/src/MXNet/NN/DataIter/Class.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/DataIter/Class.hs @@ -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 diff --git a/src/MXNet/NN/DataIter/LazyVec.hs b/src/MXNet/NN/DataIter/LazyVec.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/DataIter/LazyVec.hs @@ -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) diff --git a/src/MXNet/NN/EvalMetric.hs b/src/MXNet/NN/EvalMetric.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/EvalMetric.hs @@ -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 + diff --git a/src/MXNet/NN/Layer.hs b/src/MXNet/NN/Layer.hs --- a/src/MXNet/NN/Layer.hs +++ b/src/MXNet/NN/Layer.hs @@ -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 diff --git a/src/MXNet/NN/Optimizer.hs b/src/MXNet/NN/Optimizer.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/Optimizer.hs @@ -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 diff --git a/src/MXNet/NN/Types.hs b/src/MXNet/NN/Types.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/Types.hs @@ -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) diff --git a/src/MXNet/NN/Utils.hs b/src/MXNet/NN/Utils.hs --- a/src/MXNet/NN/Utils.hs +++ b/src/MXNet/NN/Utils.hs @@ -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