fei-nn 0.2.0 → 1.0.0
raw patch · 25 files changed
+2112/−1229 lines, 25 filesdep +FontyFruitydep +JuicyPixelsdep +Rasterificdep −attoparsecdep −attoparsec-binarydep −bytestring
Dependencies added: FontyFruity, JuicyPixels, Rasterific, conduit, conduit-combinators, formatting, neptune-backend, repa, rio, stm-conduit, streaming, type-combinators, uuid, wl-pprint-text
Dependencies removed: attoparsec, attoparsec-binary, bytestring, exceptions, fei-nn, ghc-prim, mmorph, mtl, text, unordered-containers, vector
Files
- examples/mnist/DatasetVector.hs +0/−64
- examples/mnist/Parse.hs +0/−33
- examples/mnist/lenet.hs +0/−121
- fei-nn.cabal +56/−32
- src/MXNet/NN.hs +93/−399
- src/MXNet/NN/Callback.hs +16/−27
- src/MXNet/NN/DataIter/Class.hs +19/−17
- src/MXNet/NN/DataIter/Conduit.hs +87/−0
- src/MXNet/NN/DataIter/ConduitAsync.hs +27/−0
- src/MXNet/NN/DataIter/Streaming.hs +88/−0
- src/MXNet/NN/DataIter/Vec.hs +10/−12
- src/MXNet/NN/EvalMetric.hs +194/−98
- src/MXNet/NN/Initializer.hs +36/−30
- src/MXNet/NN/Layer.hs +400/−80
- src/MXNet/NN/LrScheduler.hs +21/−21
- src/MXNet/NN/Module.hs +291/−0
- src/MXNet/NN/NDArray.hs +0/−19
- src/MXNet/NN/Optimizer.hs +123/−66
- src/MXNet/NN/Session.hs +115/−0
- src/MXNet/NN/TaggedState.hs +50/−0
- src/MXNet/NN/Types.hs +77/−82
- src/MXNet/NN/Utils.hs +120/−49
- src/MXNet/NN/Utils/GraphViz.hs +79/−79
- src/MXNet/NN/Utils/Render.hs +66/−0
- src/MXNet/NN/Utils/Repa.hs +144/−0
− examples/mnist/DatasetVector.hs
@@ -1,64 +0,0 @@-module DatasetVector where - -import MXNet.Base (Symbol, NDArray, makeNDArray, contextCPU) -import Data.Typeable -import Control.Monad.Trans.Resource (MonadResource(..), MonadThrow(..)) -import Control.Monad.IO.Class (liftIO) -import Control.Monad (liftM2) -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.Class -import MXNet.NN.DataIter.Vec -import Parse - -type SymbolF = Symbol Float -type ArrayF = NDArray Float - -loadTrainingData :: (MonadResource m, MonadThrow m) => m (DatasetVector (ArrayF, ArrayF)) -loadTrainingData = do - v1 <- batch 128 <$> sourceImages "examples/data/train-images-idx3-ubyte" - v2 <- batch 128 <$> sourceLabels "examples/data/train-labels-idx1-ubyte" - liftIO $ liftM2 zipD (mapMD cImageToNDArray v1) (mapMD cLabelToNDArray v2) - -loadTestingData :: (MonadResource m, MonadThrow m) => m (DatasetVector (ArrayF, ArrayF)) -loadTestingData = do - v1 <- batch 1 <$> sourceImages "examples/data/t10k-images-idx3-ubyte" - v2 <- batch 1 <$> sourceLabels "examples/data/t10k-labels-idx1-ubyte" - liftIO $ liftM2 zipD (mapMD cImageToNDArray v1) (mapMD cLabelToNDArray v2) - -sourceImages :: (MonadResource m, MonadThrow m) => FilePath -> m (DatasetVector Image) -sourceImages = parseFile $ do - HeaderImg n w h <- header - count n (image w h) - -sourceLabels :: (MonadResource m, MonadThrow m) => FilePath -> m (DatasetVector Label) -sourceLabels = parseFile $ do - HeaderLbl n <- header - count n label - -parseFile :: (MonadResource m, MonadThrow m) => Parser [a] -> FilePath -> m (DatasetVector a) -parseFile parser fp = do - content <- liftIO $ BS.readFile fp - case AP.parseOnly parser content of - Left msg -> throwM $ ParseError msg - Right rt -> return $ fromListD rt - -batch :: Int -> DatasetVector a -> DatasetVector (V.Vector a) -batch n (DatasetVector vec) = (DatasetVector $ walk n vec) - where - walk n vec = V.unfoldr (\v -> if V.null v then Nothing else Just (V.splitAt n v)) vec - -mapMD f (DatasetVector vec) = DatasetVector <$> V.mapM f vec - -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
@@ -1,33 +0,0 @@-module Parse where--import Data.Attoparsec.ByteString as AP-import Data.Attoparsec.Binary as AP-import qualified Data.ByteString.Internal as BS-import qualified Data.Vector.Storable as SV--type Image = SV.Vector Float-type Label = Int--data Header = HeaderImg Int Int Int- | HeaderLbl Int--header :: AP.Parser Header-header = do- mc <- AP.anyWord32be- case mc of- 0x00000803 -> do - [d1,d2,d3] <- AP.count 3 AP.anyWord32be- return $ HeaderImg (fromIntegral d1) (fromIntegral d2) (fromIntegral d3)- 0x00000801 -> do - d1 <- AP.anyWord32be- return $ HeaderLbl (fromIntegral d1)- _ -> fail "Header type not recognised"--image :: Int -> Int -> AP.Parser Image-image w h = do- BS.PS fp ofs len <- AP.take (w*h)- let vw = SV.unsafeFromForeignPtr fp ofs len- return $ SV.map ((/255) . fromIntegral) vw--label :: AP.Parser Label-label = fromIntegral <$> AP.anyWord8
− examples/mnist/lenet.hs
@@ -1,121 +0,0 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE TypeOperators #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE OverloadedLabels #-}-module Main where--import Control.Monad (forM_, void)-import qualified Data.Vector.Storable as SV-import Control.Monad.IO.Class-import Control.Monad.Trans.Resource-import System.IO (hFlush, stdout)-import qualified Data.HashMap.Strict as M--import MXNet.Base hiding (zeros)-import qualified MXNet.Base.Operators.NDArray as A-import MXNet.NN-import MXNet.NN.DataIter.Class-import qualified MXNet.NN.Utils.GraphViz as GV--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")--- 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" (#data := x .& #kernel := [5,5] .& #num_filter := 20 .& Nil)- a1 <- activation "conv1-a" (#data := v1 .& #act_type := #tanh .& Nil)- p1 <- pooling "conv1-p" (#data := a1 .& #kernel := [2,2] .& #pool_type := #max .& Nil)-- v2 <- convolution "conv2" (#data := p1 .& #kernel := [5,5] .& #num_filter := 50 .& Nil)- a2 <- activation "conv2-a" (#data := v2 .& #act_type := #tanh .& Nil)- p2 <- pooling "conv2-p" (#data := a2 .& #kernel := [2,2] .& #pool_type := #max .& Nil)-- fl <- flatten "flatten" (#data := p2 .& Nil)-- v3 <- fullyConnected "fc1" (#data := fl .& #num_hidden := 500 .& Nil)- a3 <- activation "fc1-a" (#data := v3 .& #act_type := #tanh .& Nil)-- v4 <- fullyConnected "fc2" (#data := a3 .& #num_hidden := 10 .& Nil)- a4 <- softmaxoutput "softmax" (#data := v4 .& #label := y .& Nil)- return $ Symbol a4--range :: Int -> [Int]-range = enumFromTo 1--default_initializer :: Initializer Float-default_initializer name shp@[_] = zeros name shp-default_initializer name shp@[_,_] = xavier 2.0 XavierGaussian XavierIn name shp-default_initializer name shp = normal 0.1 name shp--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- -- GV.dotPlot net GV.Png "lenet"- sess <- initialize net $ Config { - _cfg_data = M.singleton "x" [1,28,28],- _cfg_label = ["y"],- _cfg_initializers = M.empty,- _cfg_default_initializer = default_initializer,- _cfg_context = contextCPU- }- optimizer <- makeOptimizer SGD'Mom (Const 0.0002) (#momentum := 0.9 .& #wd := 0.0001 .& Nil)-- runResourceT $ train sess $ do -- trainingData <- loadTrainingData- testingData <- loadTestingData-- liftIO $ putStrLn $ "[Train] "- forM_ (range 5) $ \ind -> do- liftIO $ putStrLn $ "iteration " ++ show ind- metric <- newMetric "train" (CrossEntropy "y")- void $ forEachD_ni trainingData $ \((t,i), (x, y)) -> do- eval <- format metric- liftIO $ putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show t ++ " " ++ eval- liftIO $ hFlush stdout- fitAndEval optimizer (M.fromList [("x", x), ("y", y)]) metric- liftIO $ putStrLn ""- - liftIO $ putStrLn $ "[Test] "- result <- forEachD_ni testingData $ \((t,i), (x, y)) -> do - liftIO $ do - putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show t- hFlush stdout- [y'] <- forwardOnly (M.fromList [("x", Just x), ("y", Nothing)])- ind1 <- liftIO $ toVector y- ind2 <- liftIO $ argmax y' >>= toVector- 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 (NDArray ys) = NDArray . head <$> A.argmax (#data := ys .& #axis := Just 1 .& Nil)
fei-nn.cabal view
@@ -1,73 +1,97 @@+cabal-version: 2.2 name: fei-nn-version: 0.2.0+version: 1.0.0 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/fei-nn-license: BSD3+license: BSD-3-Clause license-file: LICENSE author: Jiasen Wu maintainer: jiasenwu@hotmail.com-copyright: Copyright: (c) 2018 Jiasen Wu+copyright: Copyright: (c) 2020 Jiasen Wu category: Machine Learning, AI build-type: Simple-cabal-version: 1.24 +Flag mxnet_geq_10600 {+ Description: MXNet >= 1.6.0+ Default: False+}++Flag neptune {+ Description: Integrate Neptune+ Default: False+}+ Library exposed-modules: MXNet.NN- MXNet.NN.NDArray MXNet.NN.Types+ MXNet.NN.Session MXNet.NN.Utils MXNet.NN.Utils.GraphViz+ MXNet.NN.Utils.Repa+ MXNet.NN.Utils.Render MXNet.NN.Layer MXNet.NN.Optimizer MXNet.NN.LrScheduler MXNet.NN.EvalMetric MXNet.NN.Initializer MXNet.NN.Callback+ MXNet.NN.TaggedState+ MXNet.NN.Module MXNet.NN.DataIter.Class MXNet.NN.DataIter.Vec+ MXNet.NN.DataIter.Streaming+ MXNet.NN.DataIter.Conduit+ MXNet.NN.DataIter.ConduitAsync other-modules: hs-source-dirs: src ghc-options: -Wall default-language: Haskell2010 default-extensions: GADTs, TypeFamilies,- OverloadedLabels+ TypeOperators,+ OverloadedLabels,+ OverloadedStrings,+ FlexibleContexts,+ FlexibleInstances,+ LambdaCase,+ MultiWayIf,+ DoAndIfThenElse,+ TypeApplications,+ DataKinds,+ RecordWildCards,+ ExplicitForAll,+ ExistentialQuantification,+ NoImplicitPrelude if impl(ghc >= 8.6) default-extensions: NoMonadFailDesugaring build-depends: base >= 4.7 && < 5.0- , unordered-containers >= 0.2.8 , resourcet >= 1.1.8- , vector >= 0.12- , mtl >= 2.2 , lens >= 4.12 , transformers-base >= 0.4.4 , aeson >= 1.2 , containers >= 0.5 , template-haskell >= 2.12 , graphviz- , text >= 1.2- , bytestring >= 0.10- , exceptions >= 0.8.3 , time < 2.0- , fei-base--Executable lenet- main-is: lenet.hs- other-modules: Parse DatasetVector- hs-source-dirs: examples/mnist- ghc-options: -Wall- default-language: Haskell2010- build-depends: base >= 4.7 && < 5.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- , ghc-prim+ , type-combinators+ , streaming >= 0.1.4.5+ , conduit >= 1.2 && < 1.4+ , conduit-combinators >= 1.1.2 && < 1.4+ , stm-conduit+ , formatting+ , wl-pprint-text+ , repa+ , Rasterific+ , JuicyPixels+ , FontyFruity+ , rio+ , uuid , fei-base- , fei-nn+ if flag(mxnet_geq_10600) {+ cpp-options: -DMXNET_GEQ_10600+ }+ if flag(neptune) {+ cpp-options: -DNEPTUNE+ build-depends: neptune-backend < 0.2+ }
src/MXNet/NN.hs view
@@ -1,422 +1,116 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE RecordWildCards #-}-{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskell #-} module MXNet.NN (- Parameter(..),- Config(..),- Session(..),- Exc(..),- Initializer,- TrainM,- CallbackClass(..), Callback(..),- train,- initialize,- fit, fit_, fitAndEval, fitDataset,- forwardOnly,- getContext,- sess_param, sess_context, sess_callbacks,+ module MXNet.NN.Module, module MXNet.NN.Optimizer, module MXNet.NN.LrScheduler, module MXNet.NN.EvalMetric,- module MXNet.NN.Initializer, module MXNet.NN.Layer,+ module MXNet.NN.Types,+ module MXNet.NN.Utils,+ module MXNet.NN.Utils.Repa,+ module MXNet.NN.Utils.GraphViz,+ module MXNet.NN.TaggedState,+ module MXNet.NN.Session, module MXNet.NN.Callback,+ module MXNet.NN.DataIter.Class,+ FeiApp,+ FeiM,+ fa_log_func,+ fa_process_context,+ fa_session,+ fa_extra,+ runFeiM,+#ifdef NEPTUNE+ runFeiM'nept,+ neptLog,+#endif+ initSession, ) where -import qualified Data.HashMap.Strict as M-import qualified Control.Monad.State.Strict as ST-import Data.Maybe (isJust, fromJust, maybe)-import Data.Foldable (forM_)-import Control.Monad (when, unless, void)-import Control.Monad.IO.Class (MonadIO, liftIO)-import Control.Monad.Trans.Resource (MonadThrow(..))-import Control.Lens (traverseOf, use, (+=), (^.), (%=))-import System.IO (hFlush, stdout)-import System.Mem-import Text.Printf-import Data.Dynamic (toDyn)--- import Control.Lens.Tuple--import MXNet.Base-import qualified MXNet.Base.Operators.NDArray as A--import MXNet.NN.NDArray-import MXNet.NN.Types-import MXNet.NN.Optimizer-import MXNet.NN.EvalMetric-import MXNet.NN.Layer-import MXNet.NN.Initializer-import MXNet.NN.LrScheduler-import MXNet.NN.DataIter.Class-import MXNet.NN.Callback---- | Execute the 'TrainM' monad-train :: (DType a, Monad m) => Session a -> TrainM a m r -> m r-train sess proc = ST.evalStateT (ST.evalStateT proc sess) (Statistics 0 0)----- | infer the shapes of input and auxiliary symbols in a symbolic neural network-inferShape' :: DType a => Symbol a -> M.HashMap String [Int] -> IO (M.HashMap String [Int], M.HashMap String [Int])-inferShape' sym known = do- (args, _, auxs, complete) <- inferShape sym (M.toList known)- unless complete $ throwM InferredShapeInComplete- return (M.fromList args, M.fromList auxs)+import Control.Lens+import Control.Monad.Trans.Resource+import RIO hiding (view)+import RIO.Process --- | initialize all parameters-initialize :: DType a => Symbol a -> Config a -> IO (Session a)-initialize sym config = do- -- give a initial batch_size = 1 for the placeholders- let spec1 = M.map (1:) $ M.difference input_shapes initializers- spec2 = initializers- dinit = config ^. cfg_default_initializer- cxt = config ^. cfg_context- (arg_with_shp, aux_with_shp) <- inferShape' sym spec1- ---------------------- -- important! labels should be merged into placeholders,- -- otherwise the labels are considered to have gradient.- ---------------------- let lbl_with_shp = M.filterWithKey (\k v -> k `elem` label_names) arg_with_shp- placeholders <- mapM (flip makeEmptyNDArray cxt) $ M.union spec1 lbl_with_shp+import MXNet.Base+import MXNet.NN.Callback+import MXNet.NN.DataIter.Class+import MXNet.NN.EvalMetric+import MXNet.NN.Layer+import MXNet.NN.LrScheduler+import MXNet.NN.Module+import MXNet.NN.Optimizer+import MXNet.NN.Session+import MXNet.NN.TaggedState+import MXNet.NN.Types+import MXNet.NN.Utils+import MXNet.NN.Utils.GraphViz+import MXNet.NN.Utils.Repa - arg_tensors <- M.traverseWithKey (initI placeholders spec2 dinit) arg_with_shp- aux_tensors <- M.traverseWithKey (initA dinit) aux_with_shp+#ifdef NEPTUNE+import Neptune.Client+import Neptune.Session (Experiment, NeptuneSession)+#endif - return $ Session {- _sess_symbol = sym,- _sess_data = input_shapes,- _sess_label = label_names,- _sess_param = arg_tensors `M.union` aux_tensors,- _sess_context = cxt,- _sess_callbacks = [],- _sess_store = M.empty+data FeiApp t n x = FeiApp+ { _fa_log_func :: !LogFunc+ , _fa_process_context :: !ProcessContext+ , _fa_session :: MVar (TaggedModuleState t n)+ , _fa_extra :: x }- where- input_shapes = config ^. cfg_data- label_names = config ^. cfg_label- initializers = config ^. cfg_initializers- -- initialize input symbols.- -- placeholders are backed by empty NDArray,- -- other input symbols are initialized by an initializer.- initI placeholder spec2 dinit inp shp =- case M.lookup inp placeholder of- Just in_arg -> do- return $ ParameterI in_arg Nothing- Nothing -> do- arg_in <- case M.lookup inp spec2 of- Just cinit -> cinit inp shp (_cfg_context config)- Nothing -> dinit inp shp (_cfg_context config)- arg_gr <- makeEmptyNDArray shp (_cfg_context config)- return $ ParameterI arg_in (Just arg_gr)- -- initialize auxiliary symbols.- initA dinit aux shp = do- arg_aux <- dinit aux shp (_cfg_context config)- return $ ParameterA arg_aux+makeLenses ''FeiApp --- | bind the symbolic network with actual parameters-bind :: (DType a, MonadIO m, MonadThrow m) => M.HashMap String (Maybe (NDArray a)) -> Bool -> TrainM a m (Executor a)-bind dat train_ = do- Context{..} <- use sess_context- net <- use sess_symbol+instance HasLogFunc (FeiApp t n x) where+ logFuncL = fa_log_func - let inputs = M.map fromJust $ M.filter isJust dat- input_shapes <- liftIO $ mapM ndshape inputs- (inp_shps, aux_shps) <- liftIO $ inferShape' net input_shapes- modifyT . traverseOf sess_param $ M.traverseWithKey $ \k p ->- case p of- ParameterI {} -> do- let ishp = inp_shps M.! k- case M.lookup k dat of- -- if the name is given in the binding data, we check its consistency.- Just a -> liftIO $ ensure_consistency (maybe (Right ishp) Left a) p- -- if the name is missing in the binding data, we check the infered shape- -- matches both the _param_in and _param_grad- Nothing -> do- pshp1 <- liftIO $ ndshape (_param_in p)- when (ishp /= pshp1 ) (throwM $ MismatchedShapeOfSym (k ++ "[i]") ishp pshp1)- case (train_, _param_grad p) of- (True, Just ndarray) -> do- pshp2 <- liftIO $ ndshape ndarray- when (ishp /= pshp2) (throwM $ MismatchedShapeOfSym (k ++ "[t]") ishp pshp2)- _ -> return ()- return p- ParameterA {} -> do- let ishp = aux_shps M.! k- pshp1 <- liftIO $ ndshape (_param_aux p)- when (ishp /= pshp1 ) (throwM $ MismatchedShapeOfSym (k ++ "[i]") ishp pshp1)- return p+instance HasSessionRef (FeiApp t n x) (TaggedModuleState t n) where+ sessionRefL = fa_session - args <- use sess_param- exec_handle <- liftIO $ do- names <- mxSymbolListArguments (unSymbol net)- -- the parameters to bind should be arranged in the same order as the names- let num_args = length names- arg_in = map (unNDArray . _param_in) $ map (args M.!) names- arg_gr = if train_- then map (fmap unNDArray . _param_grad) $ map (args M.!) names- else replicate num_args Nothing- arg_gr_req = replicate num_args (if train_ then 1 else 0)+type FeiM t n x a = ReaderT (FeiApp t n x) (ResourceT IO) a - auxnames <- mxSymbolListAuxiliaryStates (unSymbol net)- let aux_arg_aux = map (unNDArray . _param_aux) $ map (args M.!) auxnames- mxExecutorBind (unSymbol net) _device_type _device_id- arg_in arg_gr arg_gr_req- aux_arg_aux- return $ Executor exec_handle- where- -- make sure the _param_in can be used in the inference and backpropagation- -- + user data input can be in a different context w.r.t. session configuration- -- + copy inp with the right context- -- + batch size can be different from the initial configuration, or at the time- -- to swap training and inference- -- + create one and copy it- -- + for inferenceOnly, labels' NDArray can be uninitialized.- -- + just create one- ensure_consistency :: DType a => Either (NDArray a) [Int] -> Parameter a -> IO (Parameter a)- ensure_consistency (Left a) p = do- src_cxt <- context a- src_shp <- ndshape a- dst_cxt <- context (_param_in p)- dst_shp <- ndshape (_param_in p)- case (src_cxt == dst_cxt, src_shp == dst_shp) of- (True , True) -> return $ p {_param_in = a}- (False, True) -> do- A._copyto_upd [unNDArray (_param_in p)] (#data := unNDArray a .& Nil)- return p- _ -> do- a_copy <- makeEmptyNDArray src_shp dst_cxt- A._copyto_upd [unNDArray a_copy] (#data := unNDArray a .& Nil)- return $! p {_param_in = a_copy}- ensure_consistency (Right src_shp) p = do- dst_cxt <- context (_param_in p)- dst_shp <- ndshape (_param_in p)- if src_shp == dst_shp- then return p- else do- dummy <- makeEmptyNDArray src_shp dst_cxt- return $! p {_param_in = dummy} --- | single step train. Must provide all the placeholders.-fit :: (DType a, MonadIO m, MonadThrow m, Optimizer opt)- => opt a -> M.HashMap String (NDArray a) -> TrainM a m (Executor a)-fit opt datAndLbl = do- exec <- bind (M.map Just datAndLbl) True- liftIO $ do- mxExecutorForward (unExecutor exec) True- mxExecutorBackward (unExecutor exec) []- -- 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- updateParameters opt datAndLbl- return exec+data SessionAlreadyExist = SessionAlreadyExist+ deriving (Typeable, Show)+instance Exception SessionAlreadyExist --- | single step train. Must provide all the placeholders.-fit_ :: (DType a, MonadIO m, MonadThrow m, Optimizer opt)- => opt a -> M.HashMap String (NDArray a) -> TrainM a m ()-fit_ opt datAndLbl = void $ fit 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, EvalMetricMethod mtr)- => opt a -> M.HashMap String (NDArray a) -> MetricData mtr a -> TrainM a m ()-fitAndEval opt datAndLbl metric = do- Executor exec <- fit opt datAndLbl- pred <- liftIO $ map NDArray <$> mxExecutorOutputs exec- eval_results <- evaluate metric datAndLbl pred- sess_store %= M.union (M.map toDyn eval_results)--fitDataset :: (Dataset d, DatasetProp d e, DType a,- MonadIO m, MonadThrow m, DatasetConstraint d (TrainM a m),- Optimizer opt, EvalMetricMethod mtr)- => d e- -> d e- -> ([String] -> e -> M.HashMap String (NDArray a))- -> opt a- -> mtr a- -> Int- -> TrainM a m ()-fitDataset trainDataset valDataset make_binding opt metric epochs = do- callbacks <- use sess_callbacks-- data_vars <- M.keys <$> use sess_data- labl_vars <- use sess_label-- total <- sizeD trainDataset- batchSize <- batchSizeD trainDataset >>= maybe (throwM DatasetOfUnknownBatchSize) return-- liftIO $ putStrLn $ "[Train]"- forM_ (enumFromTo 1 epochs) $ \epochInd -> do- trainMetricData <- newMetric "train" metric-- liftIO $ putStrLn $ "epoch " ++ show epochInd- forM_ callbacks (begOfEpoch epochInd total)-- void $ forEachD_i trainDataset $ \(i, item) -> do- forM_ callbacks (begOfBatch i batchSize)- let binding = make_binding (data_vars ++ labl_vars) item- fitAndEval opt binding trainMetricData- eval <- format trainMetricData- liftIO $ putStr $ printf "\r\ESC[K%d/%d %s" i total eval- forM_ callbacks (endOfBatch i batchSize)- liftIO $ hFlush stdout-- forM_ callbacks (endOfEpoch epochInd total)- liftIO $ hFlush stdout- liftIO performGC-- liftIO $ putStrLn "\n[Validate]"- valMetricData <- newMetric "val" metric- void $ forEachD valDataset $ \item -> do- let whole_binding = make_binding (data_vars ++ labl_vars) item- infer_binding = M.map Just $ M.filterWithKey (const . (`elem` data_vars)) whole_binding- pred <- forwardOnly infer_binding- evaluate valMetricData whole_binding pred- eval <- format valMetricData- liftIO $ putStrLn eval-- forM_ callbacks (endOfVal epochInd total)- liftIO $ putStrLn ""--fitDataset_ :: (Dataset d, DatasetProp d e, DType a,- MonadIO m, MonadThrow m, DatasetConstraint d (TrainM a m),- Optimizer opt, EvalMetricMethod mtr)- => d e- -> ([String] -> e -> M.HashMap String (NDArray a))- -> opt a- -> MetricData mtr a- -> TrainM a m ()-fitDataset_ dataset make_binding opt metric = do- callbacks <- use sess_callbacks- total <- sizeD dataset- batchSize <- batchSizeD dataset >>= maybe (throwM DatasetOfUnknownBatchSize) return-- data_vars <- M.keys <$> use sess_data- labl_vars <- use sess_label-- void $ forEachD_i dataset $ \(i, item) -> do- forM_ callbacks (begOfBatch i batchSize)- let binding = make_binding (data_vars ++ labl_vars) item- fitAndEval opt binding metric- eval <- format metric- liftIO $ putStr $ printf "\r\ESC[K%d/%d %s" i total eval- forM_ callbacks (endOfBatch i batchSize)- liftIO $ hFlush stdout- liftIO performGC-- liftIO $ hFlush stdout--- fitDataset :: (Dataset d, DType a, DataItem e a,--- MonadIO m, MonadThrow m, DatasetConstraint d (TrainM a m),--- Optimizer opt, EvalMetricMethod mtr)--- => opt a -> Symbol a--- -> [String]--- -> d e--- -> mtr a--- -> TrainM a m ()--- fitDataset opt net varnames dataset metric = do--- callbacks <- use sess_callbacks--- shapes <- use sess_placeholders--- total <- sizeD dataset--- [example0] <- takeD 1 dataset--- batchSize <- batchSizeD example0--- -- assuming the data shape axis-0 is the batch-size--- Executor exec <- bind' net (M.map (\shp -> batchSize:tail shp) shapes) True---- forM_ callbacks (begOfEpoch total)--- void $ forEachD_i dataset $ \(i, e) -> do--- forM_ callbacks (begOfBatch i batchSize)---- t1 <- liftIO getCurrentTime---- placeHolders <- makePlaceholderMapD e varnames--- setPlaceholders placeHolders---- t2 <- liftIO getCurrentTime--- sess_prof . _1 += diffUTCTime t2 t1---- liftIO $ do--- mxExecutorForward exec True--- mxExecutorBackward exec []---- t3 <- liftIO getCurrentTime--- sess_prof . _2 += diffUTCTime t3 t2---- updateParameters opt placeHolders---- t4 <- liftIO getCurrentTime--- sess_prof . _3 += diffUTCTime t4 t3---- preds <- liftIO $ map NDArray <$> mxExecutorOutputs exec---- t5 <- liftIO getCurrentTime--- sess_prof . _4 += diffUTCTime t5 t4---- evaluate metric placeHolders preds--- eval <- format metric---- t6 <- liftIO getCurrentTime--- sess_prof . _5 += diffUTCTime t6 t5---- liftIO $ putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show total ++ " " ++ eval--- forM_ callbacks (endOfBatch i batchSize)--- liftIO $ hFlush stdout--- liftIO performGC---- t7 <- liftIO getCurrentTime--- sess_prof . _6 += diffUTCTime t7 t6--- forM_ callbacks (endOfEpoch total)--- liftIO $ hFlush stdout--updateParameters :: (MonadIO m, Optimizer opt, DType dtype)- => opt dtype -> M.HashMap String any -> TrainM dtype m ()-updateParameters opt blacklist = do- params <- use sess_param- forM_ (M.toList params) $ \ (k, v) ->- case (v, M.member k blacklist, _param_grad v) of- (ParameterI {}, False, Just grad) -> ST.lift $ optimize opt k (_param_in v) grad- _ -> return ()- ST.lift (stat_num_upd += 1)- -- waitParams---- | forward only. Must provide all the placeholders, setting the data to @Just xx@, and set label to @Nothing@.------ Note that the batch size here can be different from that in the training phase.-forwardOnly :: (DType a, MonadIO m, MonadThrow m) => M.HashMap String (Maybe (NDArray a)) -> TrainM a m [NDArray a]-forwardOnly dat = do- Executor exec <- bind dat False- liftIO $ mxExecutorForward exec False- liftIO $ map NDArray <$> mxExecutorOutputs exec+runFeiM :: x -> FeiM n t x a -> IO a+runFeiM x body = do+ -- call mxListAllOpNames can ensure the MXNet itself is properly initialized+ -- i.e. MXNet operators are registered in the NNVM+ void mxListAllOpNames+ logopt <- logOptionsHandle stdout False+ pcontx <- mkDefaultProcessContext+ session <- newEmptyMVar+ runResourceT $ do+ _ <- register mxNotifyShutdown+ withLogFunc logopt $ \logfunc ->+ flip runReaderT (FeiApp logfunc pcontx session x) body -waitParams :: (MonadIO m, DType a) => TrainM a m ()-waitParams = do- params <- use sess_param- forM_ params (\param ->- case param of- ParameterA arr1 ->- wait arr1- ParameterI arr1 Nothing ->- wait arr1- ParameterI arr1 (Just arr2) -> do- wait arr1- wait arr2- )+#ifdef NEPTUNE+type Extra'Nept x = (x, NeptuneSession, Experiment, Text -> Double -> IO ())+runFeiM'nept :: FloatDType t+ => Text -> x -> FeiM t n (Extra'Nept x) a -> IO a+runFeiM'nept project x body =+ withNept project $ \ nsess nexpt ->+ let logger k v = nlog nexpt k (fromRational (toRational v) :: Double)+ in runFeiM (x, nsess, nexpt, logger) body -wait :: (MonadIO m, DType a) => NDArray a -> TrainM a m ()-wait (NDArray hdl) = liftIO $ mxNDArrayWaitToRead hdl+neptLog :: Text -> Double -> FeiM t n (Extra'Nept x) ()+neptLog key value = do+ logger <- view $ fa_extra . _4+ liftIO $ logger key value -getContext :: Monad m => TrainM a m Context-getContext = use sess_context+#endif --- | modify the state within the inner monad------ thanks to lens, we can modify the first field of the state with following--- combinator:------ modifyT . traverseOf _1--- :: (Field1 s s a b, Monad m) => (a -> m b) -> StateT s m ()-modifyT :: Monad m => (s -> m s) -> ST.StateT s m ()-modifyT func = do- s0 <- ST.get- s1 <- ST.lift $ func s0- ST.put s1+initSession :: forall n t x. FloatDType t => SymbolHandle -> Config t -> FeiM t n x ()+initSession sym cfg = do+ sess_ref <- view $ fa_session+ liftIO $ do+ sess <- initialize sym cfg+ succ <- tryPutMVar sess_ref sess+ when (not succ) $ throwM SessionAlreadyExist
src/MXNet/NN/Callback.hs view
@@ -1,30 +1,25 @@ module MXNet.NN.Callback where -import Control.Monad.State.Strict (lift)+import RIO+import RIO.Time+import RIO.FilePath+import Formatting import Control.Lens (use)-import Text.Printf (printf)-import Control.Monad.IO.Class (liftIO)-import Control.Applicative (Alternative(..))-import Data.IORef-import Data.Dynamic (fromDynamic)-import Data.Maybe (fromMaybe)-import Data.Monoid (Alt(..))-import Data.Time.Clock (UTCTime, getCurrentTime, diffUTCTime)-import qualified Data.HashMap.Strict as M -import MXNet.NN.Types-import MXNet.NN.Utils+import MXNet.NN.Types (mod_statistics, stat_last_lr)+import MXNet.NN.Session+import MXNet.NN.TaggedState (untag)+import MXNet.NN.Utils (saveState) -- | Learning rate data DumpLearningRate = DumpLearningRate instance CallbackClass DumpLearningRate where endOfBatch _ _ _ = do- lr <- lift $ use stat_last_lr- liftIO $ do- putStr $ printf "<lr: %0.6f> " lr+ lr <- use (untag . mod_statistics . stat_last_lr)+ lift . logInfo . display $ sformat ("<lr: " % fixed 6 % ">") lr --- | Throughput +-- | Throughput data DumpThroughputEpoch = DumpThroughputEpoch { _tp_begin_time :: IORef UTCTime, _tp_end_time :: IORef UTCTime,@@ -38,13 +33,13 @@ liftIO $ getCurrentTime >>= writeIORef tt1Ref endOfEpoch _ _ (DumpThroughputEpoch _ tt2Ref _) = do liftIO $ getCurrentTime >>= writeIORef tt2Ref- endOfVal _ _ (DumpThroughputEpoch tt1Ref tt2Ref totalRef) = liftIO $ do+ endOfVal _ _ (DumpThroughputEpoch tt1Ref tt2Ref totalRef) = do tbeg <- readIORef tt1Ref tend <- readIORef tt2Ref let diff = realToFrac $ diffUTCTime tend tbeg :: Float total <- readIORef totalRef- putStr $ printf "Throughput: %d samepls/sec " (floor $ fromIntegral total / diff :: Int) writeIORef totalRef 0+ lift . logInfo . display $ sformat ("Throughput: " % int % " samples/sec") (floor $ fromIntegral total / diff :: Int) dumpThroughputEpoch :: IO Callback dumpThroughputEpoch = do@@ -55,15 +50,9 @@ return $ Callback $ DumpThroughputEpoch r0 r1 r2 -- | Checkpoint-data Checkpoint = Checkpoint String+data Checkpoint = Checkpoint FilePath instance CallbackClass Checkpoint where endOfVal i _ (Checkpoint path) = do- store <- use sess_store- let getKey key = fromMaybe (0 :: Float) $ getAlt $- (Alt $ M.lookup ("val_" ++ key) store >>= fromDynamic) <|>- (Alt $ M.lookup ("train_" ++ key) store >>= fromDynamic)- acc = getKey "acc"- loss = getKey "loss"- filename = printf "%s/epoch_%d_acc_%.2f_loss_%.2f" path i acc loss- saveSession filename+ let filename = path </> formatToString ("epoch_" % int) i+ saveState (i == 0) filename
src/MXNet/NN/DataIter/Class.hs view
@@ -3,39 +3,41 @@ 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 +import RIO +import RIO.Prelude.Types (MonadTrans) --- | Abstract Dataset type class. -- Available instances include 'LVec' and mxnet data-iters in package <https://github.com/pierric/mxnet-dataiter mxnet-dataiter> -class Dataset (d :: * -> *) where +class Dataset (d :: (* -> *) -> * -> *) where + type DatasetMonadConstraint d (m :: * -> *) :: Constraint -- | Create Dataset from `[]`. -- note that depending on the instance, it may or may not work with infinitive list. - fromListD :: [e] -> d e + fromListD :: (Monad m, DatasetMonadConstraint d m) => [e] -> d m e -- | Zip two Datasets - zipD :: d e1 -> d e2 -> d (e1, e2) + zipD :: (Monad m, DatasetMonadConstraint d m) => d m e1 -> d m e2 -> d m (e1, e2) -- | Get number of elements - sizeD :: (DatasetConstraint d m, Monad m) => d e -> m Int + sizeD :: (Monad m, DatasetMonadConstraint d m) => d m e -> m Int -- | Apply a function on each element of Dataset - forEachD :: (DatasetConstraint d m, Monad m) => d e -> (e -> m a) -> m [a] + forEachD :: (Monad m, DatasetMonadConstraint d m) => d m e -> (e -> m a) -> m [a] - -- | Apply a function on each element of Dataset together with the element's index. + -- | 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 :: (Monad m, DatasetMonadConstraint d m) => d m 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 + forEachD_ni :: (Monad m, DatasetMonadConstraint d m) => d m 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 - foldD :: (DatasetConstraint d m, Monad m) => (a -> e -> m a) -> a -> d e -> m a + foldD :: (Monad m, DatasetMonadConstraint d m) => (a -> e -> m a) -> a -> d m e -> m a - takeD :: Int -> d e -> d e + takeD :: (Monad m, DatasetMonadConstraint d m) => Int -> d m e -> d m e + -- | Lift from one monad into another + liftD :: (MonadTrans t, Monad m, DatasetMonadConstraint d m) => d m a -> d (t m) a -class DatasetProp (d :: * -> *) e where + +class Dataset d => DatasetProp (d :: (* -> *) -> * -> *) e where -- | Get the batch size of the dataset - batchSizeD :: (DatasetConstraint d m, Monad m) => d e -> m (Maybe Int) + batchSizeD :: (Monad m, DatasetMonadConstraint d m) => d m e -> m (Maybe Int)
+ src/MXNet/NN/DataIter/Conduit.hs view
@@ -0,0 +1,87 @@+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE TypeFamilies #-}+module MXNet.NN.DataIter.Conduit (+ ConduitData(..),+ Dataset(..),+ imageRecordIter_v1,+ imageRecordIter, mnistIter, csvIter, libSVMIter+) where++import Data.Conduit+import qualified Data.Conduit.Combinators as C+import qualified Data.Conduit.List as CL+import RIO+import RIO.Prelude (lift)++import MXNet.Base+import qualified MXNet.Base.DataIter as I+import MXNet.NN.DataIter.Class++data ConduitData m a = ConduitData+ { iter_batch_size :: Maybe Int+ , getConduit :: ConduitM () a m ()+ }++imageRecordIter_v1 :: (Fullfilled "_ImageRecordIter_v1" () args, DType a, MonadIO m)+ => ArgsHMap "_ImageRecordIter_v1" () args -> ConduitData m (NDArray a, NDArray a)+imageRecordIter_v1 args = ConduitData {+ getConduit = makeIter I._ImageRecordIter_v1 args,+ iter_batch_size = Just (args ! #batch_size)+}++imageRecordIter :: (Fullfilled "_ImageRecordIter" () args, DType a, MonadIO m)+ => ArgsHMap "_ImageRecordIter" () args -> ConduitData m (NDArray a, NDArray a)+imageRecordIter args = ConduitData {+ getConduit = makeIter I._ImageRecordIter args,+ iter_batch_size = Just (args ! #batch_size)+}++mnistIter :: (Fullfilled "_MNISTIter" () args, DType a, MonadIO m)+ => ArgsHMap "_MNISTIter" () args -> ConduitData m (NDArray a, NDArray a)+mnistIter args = ConduitData {+ getConduit = makeIter I._MNISTIter args,+ iter_batch_size = (args !? #batch_size) <|> Just 1+}++csvIter :: (Fullfilled "_CSVIter" () args, DType a, MonadIO m)+ => ArgsHMap "_CSVIter" () args -> ConduitData m (NDArray a, NDArray a)+csvIter args = ConduitData {+ getConduit = makeIter I._CSVIter args,+ iter_batch_size = Just (args ! #batch_size)+}++libSVMIter :: (Fullfilled "_LibSVMIter" () args, DType a, MonadIO m)+ => ArgsHMap "_LibSVMIter" () args -> ConduitData m (NDArray a, NDArray a)+libSVMIter args = ConduitData {+ getConduit = makeIter I._LibSVMIter args,+ iter_batch_size = Just (args ! #batch_size)+}++makeIter :: MonadIO m+ => (args -> IO DataIterHandle) -> args -> ConduitT i (NDArray a, NDArray a) m ()+makeIter creator args = do+ iter <- liftIO (creator args)+ let loop = do valid <- liftIO $ mxDataIterNext iter+ if valid == 0+ then liftIO (finalizeDataIterHandle iter)+ else do+ yieldM $ liftIO $ do+ dat <- mxDataIterGetData iter+ lbl <- mxDataIterGetLabel iter+ return (NDArray dat, NDArray lbl)+ loop+ loop++instance Dataset ConduitData where+ type DatasetMonadConstraint ConduitData m = ()+ fromListD = ConduitData Nothing . CL.sourceList+ zipD d1 d2 = ConduitData Nothing $ getZipSource $ (,) <$> ZipSource (getConduit d1) <*> ZipSource (getConduit d2)+ sizeD d = runConduit (getConduit d .| C.length)+ forEachD d proc = sourceToList $ getConduit d .| CL.mapM proc+ foldD proc unit d = runConduit (getConduit d .| C.foldM proc unit)+ takeD n d = d {getConduit = getConduit d .| C.take n}+ liftD d = d {getConduit = transPipe lift (getConduit d)}++instance DatasetProp ConduitData a where+ batchSizeD = return . iter_batch_size
+ src/MXNet/NN/DataIter/ConduitAsync.hs view
@@ -0,0 +1,27 @@+module MXNet.NN.DataIter.ConduitAsync where++import RIO+import qualified Data.Conduit as C+import qualified Data.Conduit.Combinators as C+import qualified Data.Conduit.List as CL+import qualified Data.Conduit.Async as CA++import MXNet.NN.DataIter.Class+import qualified MXNet.NN.DataIter.Conduit as DC+++newtype ConduitAsyncData m a = ConduitAsyncData (DC.ConduitData m a)++asyncConduit :: Maybe Int -> C.ConduitM () a m () -> ConduitAsyncData m a+asyncConduit sz cc = ConduitAsyncData (DC.ConduitData sz cc)+++instance Dataset ConduitAsyncData where+ type DatasetMonadConstraint ConduitAsyncData m = MonadUnliftIO m+ fromListD = ConduitAsyncData . fromListD+ zipD (ConduitAsyncData d1) (ConduitAsyncData d2) = ConduitAsyncData $ zipD d1 d2+ sizeD (ConduitAsyncData d) = sizeD d+ forEachD (ConduitAsyncData d) proc = CA.runCConduit $ DC.getConduit d CA.=$=& (CL.mapM proc C..| CL.consume)+ foldD proc unit (ConduitAsyncData d) = CA.runCConduit $ DC.getConduit d CA.=$=& C.foldM proc unit+ takeD n (ConduitAsyncData d) = ConduitAsyncData (takeD n d)+ liftD (ConduitAsyncData d) = ConduitAsyncData (liftD d)
+ src/MXNet/NN/DataIter/Streaming.hs view
@@ -0,0 +1,88 @@+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE TypeFamilies #-}+module MXNet.NN.DataIter.Streaming (+ StreamData(..),+ Dataset(..),+ imageRecordIter_v1,+ imageRecordIter, mnistIter, csvIter, libSVMIter+) where++import RIO+import RIO.Prelude (lift)+import Streaming+import Streaming.Prelude (Of (..), length_, toList_, yield)+import qualified Streaming.Prelude as S++import MXNet.Base+import qualified MXNet.Base.DataIter as I+import MXNet.NN.DataIter.Class++data StreamData m a = StreamData+ { iter_batch_size :: Maybe Int+ , getStream :: Stream (Of a) m ()+ }++imageRecordIter_v1 :: (Fullfilled "_ImageRecordIter_v1" () args, DType a, MonadIO m)+ => ArgsHMap "_ImageRecordIter_v1" () args -> StreamData m (NDArray a, NDArray a)+imageRecordIter_v1 args = StreamData {+ getStream = makeIter I._ImageRecordIter_v1 args,+ iter_batch_size = Just (args ! #batch_size)+}++imageRecordIter :: (Fullfilled "_ImageRecordIter" () args, DType a, MonadIO m)+ => ArgsHMap "_ImageRecordIter" () args -> StreamData m (NDArray a, NDArray a)+imageRecordIter args = StreamData {+ getStream = makeIter I._ImageRecordIter args,+ iter_batch_size = Just (args ! #batch_size)+}++mnistIter :: (Fullfilled "_MNISTIter" () args, DType a, MonadIO m)+ => ArgsHMap "_MNISTIter" () args -> StreamData m (NDArray a, NDArray a)+mnistIter args = StreamData {+ getStream = makeIter I._MNISTIter args,+ iter_batch_size = (args !? #batch_size) <|> Just 1+}++csvIter :: (Fullfilled "_CSVIter" () args, DType a, MonadIO m)+ => ArgsHMap "_CSVIter" () args -> StreamData m (NDArray a, NDArray a)+csvIter args = StreamData {+ getStream = makeIter I._CSVIter args,+ iter_batch_size = Just (args ! #batch_size)+}++libSVMIter :: (Fullfilled "_LibSVMIter" () args, DType a, MonadIO m)+ => ArgsHMap "_LibSVMIter" () args -> StreamData m (NDArray a, NDArray a)+libSVMIter args = StreamData {+ getStream = makeIter I._LibSVMIter args,+ iter_batch_size = Just (args ! #batch_size)+}++makeIter :: MonadIO m+ => (args -> IO DataIterHandle) -> args -> Stream (Of (NDArray a, NDArray a)) m ()+makeIter creator args = do+ iter <- liftIO (creator args)+ let loop = do valid <- liftIO $ mxDataIterNext iter+ if valid == 0+ then liftIO (finalizeDataIterHandle iter)+ else do+ item <- liftIO $ do+ dat <- mxDataIterGetData iter+ lbl <- mxDataIterGetLabel iter+ return (NDArray dat, NDArray lbl)+ yield item+ loop+ loop++instance Dataset StreamData where+ type DatasetMonadConstraint StreamData m = ()+ fromListD = StreamData Nothing . S.each+ zipD s1 s2 = StreamData Nothing $ S.zip (getStream s1) (getStream s2)+ sizeD = length_ . getStream+ forEachD dat proc = toList_ $ void $ S.mapM proc (getStream dat)+ foldD proc unit dat = S.foldM_ proc (return unit) return (getStream dat)+ takeD n dat = dat { getStream = S.take n (getStream dat) }+ liftD dat = dat { getStream = hoist lift (getStream dat) }++instance DatasetProp StreamData a where+ batchSizeD = return . iter_batch_size
src/MXNet/NN/DataIter/Vec.hs view
@@ -2,22 +2,19 @@ {-# LANGUAGE FlexibleInstances #-} module MXNet.NN.DataIter.Vec where -import Data.Vector (Vector)-import qualified Data.Vector as V-import Control.Monad (when)-import Control.Monad.IO.Class (MonadIO, liftIO)--- import Control.Monad.Trans.Resource (MonadThrow(..))+import RIO+import qualified RIO.NonEmpty as RNE+import qualified RIO.Vector.Boxed as V+import qualified RIO.Vector.Boxed.Partial as V (head) import MXNet.NN.DataIter.Class-import MXNet.NN.Types import MXNet.Base (NDArray, DType, ndshape) -newtype DatasetVector a = DatasetVector { _dsv_unwrap :: Vector a }+newtype DatasetVector (m :: * -> *) a = DatasetVector { _dsv_unwrap :: Vector a } -type instance DatasetConstraint DatasetVector m = MonadIO m- instance Dataset DatasetVector where+ type DatasetMonadConstraint DatasetVector m = MonadIO m fromListD = DatasetVector . V.fromList zipD v1 v2 = DatasetVector $ V.zip (_dsv_unwrap v1) (_dsv_unwrap v2) sizeD = return . V.length . _dsv_unwrap@@ -25,18 +22,19 @@ forEachD_i dat func = V.toList <$> V.forM (V.indexed $ _dsv_unwrap dat) func foldD func ele = V.foldM' func ele . _dsv_unwrap takeD n = DatasetVector . V.take n . _dsv_unwrap+ liftD (DatasetVector x) = DatasetVector x instance DType a => DatasetProp DatasetVector (NDArray a) where batchSizeD (DatasetVector dat) = liftIO $ do- batch_size : _ <- ndshape $ V.head dat+ batch_size <- RNE.head <$> ndshape (V.head dat) return $ Just batch_size instance DType a => DatasetProp DatasetVector (NDArray a, NDArray a) where batchSizeD (DatasetVector dat) = do let (arr1, arr2) = V.head dat liftIO $ do- batch_size1 : _ <- ndshape arr1- batch_size2 : _ <- ndshape arr2+ batch_size1 <- RNE.head <$> ndshape arr1+ batch_size2 <- RNE.head <$> ndshape arr2 return $ if batch_size1 /= batch_size2 then Nothing else Just batch_size1
src/MXNet/NN/EvalMetric.hs view
@@ -1,136 +1,232 @@-{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeOperators #-} module MXNet.NN.EvalMetric where -import Data.IORef --- import Data.Dynamic -import qualified Data.HashMap.Strict as M -import Control.Monad.Trans.Resource (MonadThrow(..)) -import Control.Monad -import Control.Monad.IO.Class (MonadIO, liftIO) -import Text.Printf (printf) -import qualified Data.Vector.Storable as SV +import Formatting (fixed, int, sformat, stext, (%)) +import RIO +import qualified RIO.HashMap as M +import qualified RIO.HashMap.Partial as M ((!)) +import qualified RIO.NonEmpty as RNE +import qualified RIO.Text as T +import qualified RIO.Vector.Storable as SV +import qualified RIO.Vector.Storable.Partial as SV (head) -import MXNet.Base -import qualified MXNet.Base.Operators.NDArray as A -import MXNet.NN.Types +import MXNet.Base +import MXNet.Base.Operators.Tensor (_norm) +import MXNet.NN.Layer +import MXNet.NN.Types -- | Abstract Evaluation type class class EvalMetricMethod metric where data MetricData metric a - newMetric :: (MonadIO m, DType a) - => String -- phase name - -> metric a -- tag - -> m (MetricData metric a) - evaluate :: (MonadIO m, DType a) - => MetricData metric a -- evaluation metric - -> M.HashMap String (NDArray a) -- network bindings - -> [NDArray a] -- output of the network - -> m (M.HashMap String Double) - format :: (MonadIO m, DType a) => MetricData metric a -> m String + newMetric :: (MonadIO m, FloatDType a, HasCallStack) + => Text -- phase name + -> metric a -- tag + -> m (MetricData metric a) + metricUpdate :: (MonadIO m, FloatDType a, HasCallStack) + => MetricData metric a -- evaluation metric + -> M.HashMap Text (NDArray a) -- network bindings + -> [NDArray a] -- output of the network + -> m (M.HashMap Text Double) + metricName :: MetricData metric a -> Text + metricValue :: (MonadIO m, FloatDType a, HasCallStack) => MetricData metric a -> m Double + metricFormat :: (MonadIO m, FloatDType a, HasCallStack) => MetricData metric a -> m Text + metricFormat m = do + name <- pure (metricName m) + value <- metricValue m + return $ sformat ("<" % stext % ": " % fixed 4 % ">") name value -- | Basic evaluation - accuracy -data Accuracy a = Accuracy String +data AccuracyPredType = PredByThreshold Float + | PredByArgmax + | PredByArgmaxAt Int +data Accuracy a = Accuracy + { _mtr_acc_name :: Maybe Text + , _mtr_acc_type :: AccuracyPredType + , _mtr_acc_min_value :: Float -- | to filter values less than min + , _mtr_acc_get_prob :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a + , _mtr_acc_get_gt :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a + } instance EvalMetricMethod Accuracy where - data MetricData Accuracy a = AccuracyData String String (IORef Int) (IORef Int) - newMetric phase (Accuracy label) = do + data MetricData Accuracy a = AccuracyPriv (Accuracy a) Text (IORef Int) (IORef Int) + newMetric phase conf = do a <- liftIO $ newIORef 0 b <- liftIO $ newIORef 0 - return $ AccuracyData phase label a b - evaluate (AccuracyData phase label cntRef sumRef) bindings [output] = do - liftIO $ compute output (bindings M.! label) - s <- liftIO $ readIORef sumRef - n <- liftIO $ readIORef cntRef - let acc = fromIntegral s / fromIntegral n - return $ M.singleton (phase ++ "_acc") acc - where - compute preds@(NDArray preds_hdl) lbl = do - [pred_cat_hdl] <- A.argmax (#data := preds_hdl .& #axis := Just 1 .& Nil) - pred_cat <- toVector (NDArray pred_cat_hdl) - real_cat <- toVector lbl + return $ AccuracyPriv conf phase a b - batch_size:_ <- ndshape preds - let correct = SV.length $ SV.filter id $ SV.zipWith (==) pred_cat real_cat - modifyIORef sumRef (+ correct) - modifyIORef cntRef (+ batch_size) - format (AccuracyData _ _ cntRef sumRef) = liftIO $ do + metricUpdate mtr@(AccuracyPriv Accuracy{..} phase cntRef sumRef) bindings outputs = liftIO $ do + out <- toCPU $ _mtr_acc_get_prob bindings outputs + lbl <- toCPU $ _mtr_acc_get_gt bindings outputs + + out <- case _mtr_acc_type of + PredByThreshold thr -> geqScalar thr out + PredByArgmax -> argmax out (Just (-1))False + PredByArgmaxAt axis -> argmax out (Just axis) False + + valid <- geqScalar _mtr_acc_min_value lbl + correct <- and_ valid =<< eq_ out lbl + num_correct <- SV.head <$> (toVector =<< sum_ correct Nothing False) + num_valid <- SV.head <$> (toVector =<< sum_ valid Nothing False) + + modifyIORef sumRef (+ floor num_correct) + modifyIORef cntRef (+ floor num_valid) + + value <- metricValue mtr + return $ M.singleton (metricName mtr) value + + metricName (AccuracyPriv Accuracy{..} phase _ _) = + let name = fromMaybe "acc" _mtr_acc_name + in sformat (stext % "_" % stext) phase name + + metricValue (AccuracyPriv _ _ cntRef sumRef) = liftIO $ do s <- liftIO $ readIORef sumRef n <- liftIO $ readIORef cntRef - return $ printf "<Accuracy: %0.2f>" (100 * fromIntegral s / fromIntegral n :: Float) + return (100 * fromIntegral s / fromIntegral n) --- | Basic evaluation - cross entropy -data CrossEntropy a = CrossEntropy String +-- | Basic evaluation - vector norm +data Norm a = Norm + { _mtr_norm_name :: Maybe Text + , _mtr_norm_ord :: Int + , _mtr_norm_get_array :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a + } -copyTo :: DType a => NDArray a -> NDArray a -> IO () -copyTo (NDArray dst) (NDArray src) = A._copyto_upd [dst] (#data := src .& Nil) +instance EvalMetricMethod Norm where + data MetricData Norm a = NormPriv (Norm a) Text (IORef Int) (IORef Double) + newMetric phase conf = do + a <- liftIO $ newIORef 0 + b <- liftIO $ newIORef 0 + return $ NormPriv conf phase a b + metricUpdate mtr@(NormPriv Norm{..} phase cntRef sumRef) bindings preds = liftIO $ do + array <- toCPU $ _mtr_norm_get_array bindings preds + + norm <- prim _norm (#data := array .& #ord := _mtr_norm_ord .& Nil) + norm <- SV.head <$> toVector norm + batch_size :| _ <- ndshape array + + modifyIORef' sumRef (+ realToFrac norm) + modifyIORef' cntRef (+ batch_size) + + value <- metricValue mtr + return $ M.singleton (metricName mtr) value + + metricName (NormPriv Norm{..} phase _ _) = + let lk = sformat ("_L" % int) _mtr_norm_ord + name = fromMaybe lk _mtr_norm_name + in sformat (stext % "_" % stext) phase name + + metricValue (NormPriv _ _ cntRef sumRef) = liftIO $ do + s <- readIORef sumRef + n <- readIORef cntRef + return $ realToFrac s / fromIntegral n + +-- | Basic evaluation - cross entropy +data CrossEntropy a = CrossEntropy + { _mtr_ce_name :: Maybe Text + , _mtr_ce_gt_clsid :: Bool + , _mtr_ce_get_prob :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a + , _mtr_ce_get_gt :: M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a + } + instance EvalMetricMethod CrossEntropy where - data MetricData CrossEntropy a = CrossEntropyData String String (IORef Int) (IORef Float) - newMetric phase (CrossEntropy label) = do + data MetricData CrossEntropy a = CrossEntropyPriv (CrossEntropy a) Text (IORef Int) (IORef Double) + newMetric phase conf = do a <- liftIO $ newIORef 0 b <- liftIO $ newIORef 0 - return $ CrossEntropyData phase label a b - -- | 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 (CrossEntropyData phase label cntRef sumRef) bindings [output] = do - liftIO $ compute output (bindings M.! label) - s <- liftIO $ readIORef sumRef - n <- liftIO $ readIORef cntRef - let loss = realToFrac s / fromIntegral n - return $ M.singleton (phase ++ "_loss") loss - where - compute preds lbl@(NDArray labelHandle) = do - shp1 <- ndshape preds - shp2 <- ndshape lbl - when (length shp1 /= 2 || length shp2 /= 1 || head shp1 /= head shp2) (throwM $ MismatchedShapeInEval shp1 shp2) - -- before call pick, we have to make sure preds and label - -- are in the same context - NDArray preds_may_copy <- do - c1 <- context preds - c2 <- context lbl - if c1 == c2 - then return preds - else do - preds_shap <- ndshape preds - preds_copy <- makeEmptyNDArray preds_shap c2 - copyTo preds_copy preds - return preds_copy - [predprj] <- A.pick (#data := preds_may_copy .& #index := labelHandle .& Nil) - [predlog] <- A.log (#data := predprj .& Nil) - loss <- A.sum (#data := predlog .& Nil) >>= toVector . NDArray . head - modifyIORef sumRef (+ (negate $ loss SV.! 0)) - modifyIORef cntRef (+ head shp1) - format (CrossEntropyData _ _ cntRef sumRef) = liftIO $ do - s <- liftIO $ readIORef sumRef - n <- liftIO $ readIORef cntRef - return $ printf "<CrossEntropy: %0.3f>" (realToFrac s / fromIntegral n :: Float) + return $ CrossEntropyPriv conf phase a b + metricUpdate mtr@(CrossEntropyPriv CrossEntropy{..} phase cntRef sumRef) bindings preds = liftIO $ do + prob <- toCPU $ _mtr_ce_get_prob bindings preds + gt <- toCPU $ _mtr_ce_get_gt bindings preds + + (loss, num_valid) <- + if _mtr_ce_gt_clsid + then do + -- when the gt labels are class id, + -- prob: (a, ..., b, num_classes) + -- gt: (a,..,b) + -- dim(prob) = dim(gt) + 1 + -- The last dimension serves as the prob dist. + -- We pickup the prob at the label specified class. + cls_prob <- log2_ =<< addScalar 1e-5 =<< pickI gt prob + weights <- geqScalar 0 gt + cls_prob <- mul_ cls_prob weights + nloss <- toVector =<< sum_ cls_prob Nothing False + num_valid <- toVector =<< sum_ weights Nothing False + return (negate (SV.head nloss), SV.head num_valid) + else do + -- when the gt are onehot vector + -- prob: (a, .., b, num_classes) + -- gt: (a, .., b, num_classes) + -- dim(prob) == dim(gt) + term1 <- mul_ gt =<< log2_ =<< addScalar 1e-5 prob + a <- log2_ =<< addScalar 1e-5 =<< rsubScalar 1 prob + term2 <- mul_ a =<< rsubScalar 1 gt + weights <- geqScalar 0 gt + nloss <- mul_ weights =<< add_ term1 term2 + nloss <- toVector =<< sum_ nloss Nothing False + num_valid <- toVector =<< sum_ weights Nothing False + return (negate (SV.head nloss), SV.head num_valid) + + modifyIORef' sumRef (+ realToFrac loss) + modifyIORef' cntRef (+ floor num_valid) + + value <- metricValue mtr + return $ M.singleton (metricName mtr) value + + metricName (CrossEntropyPriv CrossEntropy{..} phase _ _) = + let name = fromMaybe "ce" _mtr_ce_name + in sformat (stext % "_" % stext) phase name + + metricValue (CrossEntropyPriv _ _ cntRef sumRef) = liftIO $ do + s <- readIORef sumRef + n <- readIORef cntRef + return $ realToFrac s / fromIntegral n + + data ListOfMetric ms a where MNil :: ListOfMetric '[] a (:*) :: (EvalMetricMethod m) => m a -> ListOfMetric ms a -> ListOfMetric (m ': ms) a instance EvalMetricMethod (ListOfMetric '[]) where data MetricData (ListOfMetric '[]) a = MNilData - newMetric _ _ = return MNilData - evaluate _ _ _ = return M.empty - format _ = return "" + newMetric _ _ = return MNilData + metricName _ = error "Empty metric" + metricValue _ = error "Empty metric" + metricUpdate _ _ _ = return M.empty + metricFormat _ = return "" instance (EvalMetricMethod m, EvalMetricMethod (ListOfMetric ms)) => EvalMetricMethod (ListOfMetric (m ': ms)) where data MetricData (ListOfMetric (m ': ms)) a = MCompositeData (MetricData m a) (MetricData (ListOfMetric ms) a) newMetric phase (a :* as) = MCompositeData <$> (newMetric phase a) <*> (newMetric phase as) - evaluate (MCompositeData a as) bindings output = do - m1 <- evaluate a bindings output - m2 <- evaluate as bindings output + metricUpdate (MCompositeData a as) bindings output = do + m1 <- metricUpdate a bindings output + m2 <- metricUpdate as bindings output return $ M.union m1 m2 - format (MCompositeData a as) = do - s1 <- format a - s2 <- format as - return $ s1 ++ " " ++ s2 + metricName _ = error "List of metrics" + metricValue _ = error "List of metrics" + metricFormat (MCompositeData a as) = do + s1 <- metricFormat a + s2 <- metricFormat as + return $ T.concat [s1, " ", s2] infixr 9 :* + +class MetricsToList ms where + metricsToList :: (MonadIO m, FloatDType a) => MetricData (ListOfMetric ms) a -> m [(Text, Double)] + +instance MetricsToList '[] where + metricsToList MNilData = return [] + +instance (EvalMetricMethod m, MetricsToList n) => MetricsToList (m ': n) where + metricsToList (MCompositeData a b) = do + n <- pure $ metricName a + v <- metricValue a + w <- metricsToList b + return $ (n, v) : w
src/MXNet/NN/Initializer.hs view
@@ -1,17 +1,16 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE TypeApplications #-} {-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE QuasiQuotes #-} module MXNet.NN.Initializer where -import Control.Monad.Trans.Resource (MonadThrow(..))--import MXNet.Base-import qualified MXNet.Base.Operators.NDArray as A-import qualified Data.Vector.Storable as SV+import RIO+import qualified RIO.NonEmpty as RNE+import qualified RIO.Text as T+import qualified RIO.Vector.Storable as SV -import MXNet.NN.Types-import MXNet.NN.Utils+import MXNet.Base+import qualified MXNet.Base.Operators.Tensor as T+import MXNet.Base.Tensor (prim)+import MXNet.NN.Types+import MXNet.NN.Utils empty :: DType a => Initializer a empty _ shp cxt = makeEmptyNDArray shp cxt@@ -25,37 +24,44 @@ constant :: DType a => a -> Initializer a constant val _ shp cxt = makeNDArray shp cxt $ SV.replicate (product shp) val -uniform :: forall a. (DType a, HasEnum (DTypeName a) '["None", "float16" ,"float32", "float64"]) +uniform :: forall a. (DType a, HasEnum (DTypeName a) '["None", "float16" ,"float32", "float64"]) => Float -> Initializer a-uniform sca _ shp cxt = NDArray . head <$> (A._random_uniform - ( #low := (-sca) +uniform sca _ shp cxt = prim T.__random_uniform+ ( #low := (-sca) .& #high := sca- .& #shape := shp+ .& #shape := RNE.toList shp .& #ctx := formatContext cxt .& #dtype := EnumType (typename (undefined :: a))- .& Nil))+ .& Nil) -normal :: forall a. (DType a, HasEnum (DTypeName a) '["None", "float16" ,"float32", "float64"]) +normal :: forall a. (DType a, HasEnum (DTypeName a) '["None", "float16" ,"float32", "float64"]) => Float -> Initializer a-normal sigma _ shp cxt = NDArray . head <$> (A._random_normal+normal sigma _ shp cxt = prim T.__random_normal ( #loc := (0 :: Float) .& #scale := sigma- .& #shape := shp+ .& #shape := RNE.toList shp .& #ctx := formatContext cxt .& #dtype := EnumType (typename (undefined :: a))- .& Nil))+ .& Nil) -data XavierFactor = XavierAvg | XavierIn | XavierOut-data XavierRandom = XavierUniform | XavierGaussian+data XavierFactor = XavierAvg+ | XavierIn+ | XavierOut+data XavierRandom = XavierUniform+ | XavierGaussian xavier :: (DType a, HasEnum (DTypeName a) '["None", "float16" ,"float32", "float64"]) => Float -> XavierRandom -> XavierFactor -> Initializer a-xavier magnitude distr factor name (shp@[ofan,ifan]) cxt =- let scale = case factor of - XavierIn -> sqrt (magnitude / fromIntegral ifan)- XavierOut -> sqrt (magnitude / fromIntegral ofan)- XavierAvg -> sqrt (magnitude * 2.0 / fromIntegral (ifan + ofan))- in case distr of- XavierUniform -> uniform scale name shp cxt- XavierGaussian-> normal scale name shp cxt-xavier _ _ _ _ shp _ = throwM $ InvalidArgument $ "invalid shape " ++ show shp ++ " for xavier initializer"+xavier magnitude distr factor name shp cxt+ | RNE.length shp < 2 = throwM $ InvalidArgument $+ T.concat ["invalid shape ", formatShape shp, " for xavier initializer"]+ | otherwise =+ let ofan :| dims = shp+ ifan = product dims+ scale = case factor of+ XavierIn -> sqrt (magnitude / fromIntegral ifan)+ XavierOut -> sqrt (magnitude / fromIntegral ofan)+ XavierAvg -> sqrt (magnitude * 2.0 / fromIntegral (ifan + ofan))+ in case distr of+ XavierUniform -> uniform scale name shp cxt+ XavierGaussian-> normal scale name shp cxt
src/MXNet/NN/Layer.hs view
@@ -1,100 +1,420 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE PartialTypeSignatures #-}+{-# LANGUAGE TypeFamilyDependencies #-}+{-# LANGUAGE UndecidableInstances #-}+module MXNet.NN.Layer where -module MXNet.NN.Layer (- variable,- convolution,- fullyConnected,- pooling,- activation,- softmaxoutput,- batchnorm,- cast,- plus,- flatten,- identity,- dropout,- reshape,-) where+import qualified Data.UUID as UUID+import qualified Data.UUID.V4 as UUID+import Formatting (formatToString, int, shown, stext,+ (%))+import RIO+import qualified RIO.NonEmpty as NE+import qualified RIO.State as ST+import qualified RIO.Text as RT+import System.IO.Unsafe (unsafePerformIO) -import MXNet.Base-import qualified MXNet.Base.Operators.Symbol as S+import MXNet.Base+import qualified MXNet.Base.Operators.Tensor as S+import qualified MXNet.NN.Types as S -variable :: String -> IO SymbolHandle-variable = mxSymbolCreateVariable+runLayerBuilder :: MonadIO m => Layer a -> m a+runLayerBuilder = liftIO . flip ST.evalStateT [] -convolution :: (HasArgs "_Convolution(symbol)" args '["kernel", "num_filter", "data", "stride", "dilate", "pad", "num_group", "workspace", "layout", "cudnn_tune", "cudnn_off", "no_bias"]- ,WithoutArgs "_Convolution(symbol)" args '["bias", "weight"])- => String -> ArgsHMap "_Convolution(symbol)" args -> IO SymbolHandle-convolution name args = do- b <- variable (name ++ "_bias")- w <- variable (name ++ "_weight")- if args !? #no_bias == Just True +type instance TensorMonad SymbolHandle = Layer++instance PrimTensorOp SymbolHandle SymbolHandle where+ prim op args = getNextNamePrefixed >>= liftIO . op args++type Layer = ST.StateT [(Maybe Text, SomeNameBuilder)] IO++class Show nb => NameBuilder nb where+ nextName :: MonadIO m => nb -> m Text++data SomeNameBuilder = forall nb . (NameBuilder nb) => SomeNameBuilder nb++instance Show SomeNameBuilder where+ show (SomeNameBuilder nb) = show nb++instance NameBuilder SomeNameBuilder where+ nextName (SomeNameBuilder nb) = nextName nb+++data UUIDNameBuilder = UUIDNameBuilder++instance Show UUIDNameBuilder where+ show _ = "UUID"++instance NameBuilder UUIDNameBuilder where+ nextName _ = do+ uuid <- liftIO $ UUID.nextRandom+ return $ UUID.toText uuid++newtype SequNameBuilder = SequNameBuilder (IORef Int)++instance Show SequNameBuilder where+ show (SequNameBuilder ref) =+ let idx = unsafePerformIO (readIORef ref)+ in formatToString ("Seq:" % int) idx++instance NameBuilder SequNameBuilder where+ nextName (SequNameBuilder ref) = do+ n <- liftIO $ readIORef ref+ liftIO $ writeIORef ref (n+1)+ return (tshow n)++data OnceNameBuilder = OnceNameBuilder Text (IORef Bool)++instance Show OnceNameBuilder where+ show (OnceNameBuilder n ref) =+ let flag = unsafePerformIO (readIORef ref)+ in formatToString ("Once:" % stext % "[" % shown % "]") n flag++instance NameBuilder OnceNameBuilder where+ nextName (OnceNameBuilder name flag) = do+ fresh <- readIORef flag+ if fresh+ then do+ writeIORef flag False+ return name+ else throwString (formatToString ("name \"" % stext % "\" has been used.") name)++dumpCurrentScope :: Layer Text+dumpCurrentScope = do+ scopes <- ST.get+ return $ tshow scopes++sequential :: HasCallStack => Text -> Layer a -> Layer a+sequential name mk = do+ nb <- liftIO $ SequNameBuilder <$> newIORef 0+ subscope (Just name, SomeNameBuilder nb) mk++sequential' :: HasCallStack => Layer a -> Layer a+sequential' mk = do+ nb <- liftIO $ SequNameBuilder <$> newIORef 0+ subscope (Nothing, SomeNameBuilder nb) mk++unique :: HasCallStack => Text -> Layer a -> Layer a+unique name = subscope (Just name, SomeNameBuilder UUIDNameBuilder)++unique' :: HasCallStack => Layer a -> Layer a+unique' = subscope (Nothing, SomeNameBuilder UUIDNameBuilder)++named :: HasCallStack => Text -> Layer a -> Layer a+named name mk = do+ scopes <- ST.get+ fresh <- newIORef True+ ST.put ((Nothing, SomeNameBuilder (OnceNameBuilder name fresh)) : scopes)+ a <- mk+ ST.put scopes+ return a++getNextName :: HasCallStack => Layer Text+getNextName = do+ scopes <- ST.get+ case scopes of+ ((_, nb) : _) -> nextName nb+ _ -> throwString ("No next name avaiable. The current scopes: " ++ show scopes)++getNextNamePrefixed :: HasCallStack => Layer Text+getNextNamePrefixed = do+ name <- getNextName+ getNamePrefixed (Just name)++getNamePrefixed :: HasCallStack => Maybe Text -> Layer Text+getNamePrefixed name = do+ scopes <- ST.get+ let comps = catMaybes $ reverse (map fst scopes) ++ [name]+ return $ RT.intercalate "." comps++subscope :: HasCallStack => (Maybe Text, SomeNameBuilder) -> Layer a -> Layer a+subscope scope mk = do+ old_scopes <- ST.get+ ST.put (scope : old_scopes)+ a <- mk+ ST.put old_scopes+ return a++subscope_named :: HasCallStack => Text -> Layer a -> Layer a+subscope_named name = subscope (Just name, SomeNameBuilder UUIDNameBuilder)++subscope_next_name :: HasCallStack => Layer a -> Layer a+subscope_next_name mk = do+ name <- getNextName+ subscope_named name mk+++variable :: Text -> Layer SymbolHandle+variable name = getNamePrefixed (Just name) >>= liftIO . mxSymbolCreateVariable++constant :: NonEmpty Int -> [Float] -> Layer SymbolHandle+constant shape value = do+ name <- getNextNamePrefixed+ let build = do var <- mxSymbolCreateVariable name+ mxSymbolSetAttr var "__shape__" (tshow $ NE.toList shape)+ mxSymbolSetAttr var "__init__" (tshow value)+ return var+ named (RT.concat [name, ".sg"]) $ blockGrad =<< liftIO build++convolution :: (HasArgs "_Convolution" SymbolHandle args+ '["kernel", "num_filter", "data", "stride", "dilate", "pad",+ "num_group", "workspace", "layout",+ "cudnn_tune", "cudnn_off", "no_bias"]+ ,WithoutArgs "_Convolution" SymbolHandle args+ '["bias", "weight"])+ => ArgsHMap "_Convolution" _ args -> Layer SymbolHandle+convolution args = subscope_next_name $ do+ b <- variable "bias"+ w <- variable "weight"++ name <- getNamePrefixed Nothing+ if args !? #no_bias == Just True then- S._Convolution name (#weight := w .& args)+ liftIO $ S._Convolution (#weight := w .& args) name else- S._Convolution name (#bias := b .& #weight := w .& args)+ liftIO $ S._Convolution (#bias := b .& #weight := w .& args) name -fullyConnected :: (HasArgs "_FullyConnected(symbol)" args '["flatten", "no_bias", "data", "num_hidden"]- ,WithoutArgs "_FullyConnected(symbol)" args '["bias", "weight"])- => String -> ArgsHMap "_FullyConnected(symbol)" args -> IO SymbolHandle-fullyConnected name args = do- b <- variable (name ++ "_bias")- w <- variable (name ++ "_weight")- if args !? #no_bias == Just True +convolutionShared :: (HasArgs "_Convolution" SymbolHandle args+ '["kernel", "num_filter", "stride",+ "dilate", "pad", "num_group", "workspace",+ "layout", "cudnn_tune", "cudnn_off", "no_bias"]+ ,WithoutArgs "_Convolution" SymbolHandle args+ '["data", "bias", "weight"])+ => ArgsHMap "_Convolution" _ args -> Layer (SymbolHandle -> Layer SymbolHandle)+convolutionShared args = subscope_next_name $ do+ b <- variable "bias"+ w <- variable "weight"++ return $ \data_ -> do+ name <- getNextNamePrefixed+ if args !? #no_bias == Just True+ then+ liftIO $ S._Convolution (#data := data_ .& #weight := w .& args) name+ else+ liftIO $ S._Convolution (#data := data_ .& #bias := b .& #weight := w .& args) name++fullyConnected :: (HasArgs "_FullyConnected" SymbolHandle args+ '["flatten", "no_bias", "data", "num_hidden"]+ ,WithoutArgs "_FullyConnected" SymbolHandle args+ '["bias", "weight"])+ => ArgsHMap "_FullyConnected" _ args -> Layer SymbolHandle+fullyConnected args = subscope_next_name $ do+ b <- variable "bias"+ w <- variable "weight"++ name <- getNamePrefixed Nothing+ if args !? #no_bias == Just True then- S._FullyConnected name (#weight := w .& args)+ liftIO $ S._FullyConnected (#weight := w .& args) name else- S._FullyConnected name (#bias := b .& #weight := w .& args)+ liftIO $ S._FullyConnected (#bias := b .& #weight := w .& args) name --- 1.0.0 pooling :: HasArgs "_Pooling(symbol)" args '["data", "kernel", "pool_type", "stride", "pad", "pooling_convention", "global_pool", "cudnn_off"]--- 1.4.0 pooling :: HasArgs "_Pooling(symbol)" args '["data", "kernel", "pool_type", "stride", "pad", "pooling_convention", "global_pool", "cudnn_off", "p_value", "count_include_pad"]--- 1.5.0-pooling :: HasArgs "_Pooling(symbol)" args '["data", "kernel", "pool_type", "stride", "pad", "pooling_convention", "global_pool", "cudnn_off", "p_value", "count_include_pad", "layout"]- => String -> ArgsHMap "_Pooling(symbol)" args -> IO SymbolHandle-pooling = S._Pooling+fullyConnectedShared :: (HasArgs "_FullyConnected" SymbolHandle args+ '["flatten", "no_bias", "num_hidden"]+ ,WithoutArgs "_FullyConnected" SymbolHandle args+ '["bias", "weight"])+ => ArgsHMap "_FullyConnected" _ args -> Layer (SymbolHandle -> Layer SymbolHandle)+fullyConnectedShared args = subscope_next_name $ do+ b <- variable "bias"+ w <- variable "weight" -activation :: HasArgs "_Activation(symbol)" args '["data", "act_type"]- => String -> ArgsHMap "_Activation(symbol)" args -> IO SymbolHandle-activation = S._Activation+ return $ \data_ -> do+ name <- getNextNamePrefixed+ if args !? #no_bias == Just True+ then+ liftIO $ S._FullyConnected (#data := data_ .& #weight := w .& args) name+ else+ liftIO $ S._FullyConnected (#data := data_ .& #bias := b .& #weight := w .& args) name -softmaxoutput :: HasArgs "_SoftmaxOutput(symbol)" args '["data", "label", "out_grad", "smooth_alpha", "normalization", "preserve_shape", "multi_output", "use_ignore", "ignore_label", "grad_scale"]- => String -> ArgsHMap "_SoftmaxOutput(symbol)" args -> IO SymbolHandle-softmaxoutput = S._SoftmaxOutput+batchnorm :: HasArgs "_BatchNorm" SymbolHandle args+ '["data", "eps", "momentum", "fix_gamma",+ "use_global_stats", "output_mean_var", "axis",+ "cudnn_off", "min_calib_range", "max_calib_range"]+ => ArgsHMap "_BatchNorm" _ args -> Layer SymbolHandle+batchnorm args = subscope_next_name $ do+ gamma <- variable "gamma"+ beta <- variable "beta"+ mov_mean <- variable "running_mean"+ mov_var <- variable "running_var" -batchnorm :: HasArgs "_BatchNorm(symbol)" args '["data", "eps", "momentum", "fix_gamma", "use_global_stats", "output_mean_var", "axis", "cudnn_off"]- => String -> ArgsHMap "_BatchNorm(symbol)" args -> IO SymbolHandle-batchnorm name args = do- gamma <- variable (name ++ "_gamma")- beta <- variable (name ++ "_beta")- mov_mean <- variable (name ++ "_moving_mean")- mov_var <- variable (name ++ "_moving_var")- S._BatchNorm name (#gamma := gamma .& #beta := beta .& #moving_mean := mov_mean .& #moving_var := mov_var .& args)+ name <- getNamePrefixed Nothing+ liftIO $ S._BatchNorm (#gamma := gamma+ .& #beta := beta+ .& #moving_mean := mov_mean+ .& #moving_var := mov_var+ .& args) name -cast :: HasArgs "_Cast(symbol)" args '["data", "dtype"]- => String -> ArgsHMap "_Cast(symbol)" args -> IO SymbolHandle-cast name args = S._Cast name args+blockGrad :: SymbolHandle -> Layer SymbolHandle+blockGrad s = prim S._BlockGrad (#data := s .& Nil) -plus :: HasArgs "elemwise_add(symbol)" args '["lhs", "rhs"]- => String -> ArgsHMap "elemwise_add(symbol)" args -> IO SymbolHandle-plus = S.elemwise_add+splitBySections :: HasCallStack => Int -> Int -> Bool -> SymbolHandle -> Layer [SymbolHandle]+splitBySections num_sections axis squeeze s = do+ r <- prim S.__split_v2 (#data := s+ .& #axis := axis+ .& #indices := []+ .& #sections := num_sections+ .& #squeeze_axis := squeeze .& Nil)+ mapM (at r) ([0..num_sections-1] :: [Int]) -flatten :: HasArgs "_Flatten(symbol)" args '["data"]- => String -> ArgsHMap "_Flatten(symbol)" args -> IO SymbolHandle-flatten = S._Flatten+-----------------------------------------------------------------------------+-- For both Symbol and NDArray+----------------------------------------------------------------------------- -identity :: HasArgs "_copy(symbol)" args '["data"]- => String -> ArgsHMap "_copy(symbol)" args -> IO SymbolHandle-identity = S._copy+pooling :: (PrimTensorOp t t, Fullfilled "_Pooling" t args)+ => ArgsHMap "_Pooling" t args -> TensorM t+pooling = prim S._Pooling --- 1.4.0 dropout :: HasArgs "_Dropout(symbol)" args '["data", "mode", "p", "axes"] --- 1.5.0-dropout :: HasArgs "_Dropout(symbol)" args '["data", "mode", "p", "axes", "cudnn_off"] - => String -> ArgsHMap "_Dropout(symbol)" args -> IO SymbolHandle-dropout = S._Dropout+activation :: (PrimTensorOp t t, Fullfilled "_Activation" t args)+ => ArgsHMap "_Activation" t args -> TensorM t+activation = prim S._Activation -reshape :: (HasArgs "_Reshape(symbol)" args '["data", "shape", "reverse"]- ,WithoutArgs "_Reshape(symbol)" args '["target_shape", "keep_highest"])- => String -> ArgsHMap "_Reshape(symbol)" args -> IO SymbolHandle-reshape = S._Reshape+softmax :: (PrimTensorOp t t, Fullfilled "_softmax" t args)+ => ArgsHMap "_softmax" t args -> TensorM t+softmax = prim S._softmax++softmaxoutput :: (PrimTensorOp t t, Fullfilled "_SoftmaxOutput" t args)+ => ArgsHMap "_SoftmaxOutput" t args -> TensorM t+softmaxoutput = prim S._SoftmaxOutput++pick :: (PrimTensorOp t t, Fullfilled "_pick" t args)+ => ArgsHMap "_pick" t args -> TensorM t+pick = prim S._pick++stack axis ts = prim S._stack (#num_args := length ts .& #data := ts .& #axis := axis .& Nil)+flatten t = prim S._Flatten (#data := t .& Nil)+identity s = prim S.__copy (#data := s .& Nil)+dropout t p = prim S._Dropout (#data := t .& #p := p .& Nil)+reshape shape a = prim S._Reshape (#data := a .& #shape := shape .& Nil)++add_, sub_, mul_, div_, eq_, neq_, lt_, leq_, gt_, geq_ ::+ PrimTensorOp t t => t -> t -> TensorM t+add_ a b = prim S._elemwise_add (#lhs := a .& #rhs := b .& Nil)+sub_ a b = prim S._elemwise_sub (#lhs := a .& #rhs := b .& Nil)+mul_ a b = prim S._elemwise_mul (#lhs := a .& #rhs := b .& Nil)+div_ a b = prim S._elemwise_div (#lhs := a .& #rhs := b .& Nil)++eq_ a b = prim S.__equal (#lhs := a .& #rhs := b .& Nil)+neq_ a b = prim S.__not_equal (#lhs := a .& #rhs := b .& Nil)+lt_ a b = prim S.__lesser (#lhs := a .& #rhs := b .& Nil)+leq_ a b = prim S.__lesser_equal (#lhs := a .& #rhs := b .& Nil)+gt_ a b = prim S.__greater (#lhs := a .& #rhs := b .& Nil)+geq_ a b = prim S.__greater_equal (#lhs := a .& #rhs := b .& Nil)++and_ a b = prim S.__logical_and (#lhs := a .& #rhs := b .& Nil)+or_ a b = prim S.__logical_or (#lhs := a .& #rhs := b .& Nil)+xor_ a b = prim S.__logical_xor (#lhs := a .& #rhs := b .& Nil)+not_ a = prim S._logical_not (#data := a .& Nil)++addScalar b a = prim S.__plus_scalar (#data := a .& #scalar := b .& Nil)+subScalar b a = prim S.__minus_scalar (#data := a .& #scalar := b .& Nil)+rsubScalar b a = prim S.__rminus_scalar (#data := a .& #scalar := b .& Nil)+mulScalar b a = prim S.__mul_scalar (#data := a .& #scalar := b .& Nil)+divScalar b a = prim S.__div_scalar (#data := a .& #scalar := b .& Nil)+rdivScalar b a = prim S.__rdiv_scalar (#data := a .& #scalar := b .& Nil)++eqScalar b a = prim S.__equal_scalar (#data := a .& #scalar := b .& Nil)+neqScalar b a = prim S.__not_equal_scalar (#data := a .& #scalar := b .& Nil)+ltScalar b a = prim S.__lesser_scalar (#data := a .& #scalar := b .& Nil)+leqScalar b a = prim S.__lesser_equal_scalar (#data := a .& #scalar := b .& Nil)+gtScalar b a = prim S.__greater_scalar (#data := a .& #scalar := b .& Nil)+geqScalar b a = prim S.__greater_equal_scalar (#data := a .& #scalar := b .& Nil)++andScalar b a = prim S.__logical_and_scalar (#data := a .& #scalar := b .& Nil)+orScalar b a = prim S.__logical_or_scalar (#data := a .& #scalar := b .& Nil)+xorScalar b a = prim S.__logical_xor_scalar (#data := a .& #scalar := b .& Nil)++addBroadcast a b = prim S._broadcast_add (#lhs := a .& #rhs := b .& Nil)+subBroadcast a b = prim S._broadcast_sub (#lhs := a .& #rhs := b .& Nil)+mulBroadcast a b = prim S._broadcast_mul (#lhs := a .& #rhs := b .& Nil)+divBroadcast a b = prim S._broadcast_div (#lhs := a .& #rhs := b .& Nil)++eqBroadcast a b = prim S._broadcast_equal (#lhs := a .& #rhs := b .& Nil)+neqBroadcast a b = prim S._broadcast_not_equal (#lhs := a .& #rhs := b .& Nil)+ltBroadcast a b = prim S._broadcast_lesser (#lhs := a .& #rhs := b .& Nil)+leqBroadcast a b = prim S._broadcast_lesser_equal (#lhs := a .& #rhs := b .& Nil)+gtBroadcast a b = prim S._broadcast_greater (#lhs := a .& #rhs := b .& Nil)+geqBroadcast a b = prim S._broadcast_greater_equal (#lhs := a .& #rhs := b .& Nil)++ceil_ a = prim S._ceil (#data := a .& Nil)+floor_ a = prim S._floor (#data := a .& Nil)+sqrt_ a = prim S._sqrt (#data := a .& Nil)+log2_ a = prim S._log2 (#data := a .& Nil)+square_ a = prim S._square (#data := a .& Nil)++concat_ :: PrimTensorOp t t => Int -> [t] -> TensorM t+concat_ d s = prim S._Concat (#data := s .& #num_args := length s .& #dim := d .& Nil)++takeI :: (HasCallStack, PrimTensorOp t t)+ => t -> t -> TensorM t+takeI i a = prim S._take (#a := a .& #indices := i .& Nil)++pickI :: (HasCallStack, PrimTensorOp t t)+ => t -> t -> TensorM t+pickI i t = prim S._pick (#data := t .& #index := i .& Nil)++where_ c a b = prim S._where (#condition := c .& #x := a .& #y := b .& Nil)++zerosLike a = prim S._zeros_like (#data := a .& Nil)+onesLike a = prim S._ones_like (#data := a .& Nil)++squeeze axis a = prim S._squeeze (#data := a .& #axis := axis .& Nil)+expandDims axis a = prim S._expand_dims (#data := a .& #axis := axis .& Nil)++broadcastAxis axis size a = prim S._broadcast_axis (#data := a .& #axis := axis .& #size := size .& Nil)++sum_ s axis keepdims = prim S._sum (#data := s .& #axis:= axis .& #keepdims := keepdims .& Nil)++transpose a axes = prim S._transpose (#data := a .& #axes := axes .& Nil)++argmax a axis keepdims = prim S._argmax (#data := a .& #axis := axis .& #keepdims := keepdims .& Nil)++slice_axis a axis beg end = prim S._slice_axis (#data := a .& #axis := axis .& #begin := beg .& #end := end .& Nil)++-- TODO constraint the `o` to conform to `dt`+cast :: PrimTensorOp t o+ => EnumType '["bool", "float16", "float32", "float64", "int32", "int64", "int8", "uint8"]+ -> t+ -> TensorM o+cast dt t = prim S._Cast (#dtype := dt .& #data := t .& Nil)++----------------------------------------------------------------------------+data LossAgg = AggMean | AggSum++sigmoidBCE :: (PrimTensorOp t t, Monad (TensorMonad t))+ => t -> t -> Maybe t -> LossAgg -> TensorM t+sigmoidBCE pred label sample_weight agg = do+ -- pred: (B, C, 1)+ -- label: (B, C, 1)++ a <- prim S._relu (#data := pred .& Nil)+ b <- mul_ pred label+ c <- prim S._abs (#data := pred .& Nil) >>= rsubScalar 0+ c <- prim S._Activation (#data := c .& #act_type := #softrelu .& Nil)+ loss <- add_ c =<< sub_ a b+ loss <- case sample_weight of+ Just w -> mulBroadcast loss w+ Nothing -> return loss+ case agg of+ AggMean -> prim S._mean (#data := loss .& #axis := Just [0] .& #exclude := True .& Nil)+ AggSum -> prim S._sum (#data := loss .& #axis := Just [0] .& #exclude := True .& Nil)++softmaxCE :: (PrimTensorOp t t, Monad (TensorMonad t)) => Int -> t -> t -> Maybe t -> TensorM t+softmaxCE axis pred label sample_weight = do+ pred <- prim S._log_softmax (#data := pred .& #axis := axis .& Nil)+ labl <- prim S._reshape_like (#lhs := label .& #rhs := pred .& Nil)+ loss <- mul_ pred labl+ loss <- sum_ loss (Just [axis]) True >>= rsubScalar 0+ loss <- case sample_weight of+ Just w -> mulBroadcast loss w+ Nothing -> return loss+ prim S._mean (#data := loss .& #axis := Just [0] .& #exclude := True .& Nil)++-----------------------------------------------------------------------------+-- For NDArray Only+-----------------------------------------------------------------------------+copy :: (HasCallStack, PrimTensorOp t t, TensorApply t ~ (Maybe [t] -> IO [t]))+ => t -> t -> IO t+copy src dst = do+ [ret] <- S.__copyto (#data := src .& Nil) (Just [dst])+ return ret
src/MXNet/NN/LrScheduler.hs view
@@ -1,10 +1,10 @@+{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE TypeOperators #-} module MXNet.NN.LrScheduler where -import MXNet.Base.Spec.Operator-import Data.Maybe (fromMaybe)+import MXNet.Base.Spec.Operator+import RIO hiding (Const) class LrScheduler sch where getLR :: sch -> Int -> Float@@ -21,18 +21,18 @@ data FactorScheduler = Factor Float Float Int Float instance LrScheduler FactorScheduler where- getLR (Factor base factor step stop) nup = + getLR (Factor base factor step stop) nup = let lr = base * factor ^ (nup `div` step) in if lr < stop then stop else lr -type instance ParameterList "lrOfFactor" =- '[ '("factor", 'AttrReq Float), '("step", 'AttrReq Int), +type instance ParameterList "lrOfFactor" () =+ '[ '("factor", 'AttrReq Float), '("step", 'AttrReq Int), '("base", 'AttrOpt Float), '("stop", 'AttrOpt Float)]- -lrOfFactor :: Fullfilled "lrOfFactor" args - => ArgsHMap "lrOfFactor" args -> FactorScheduler++lrOfFactor :: Fullfilled "lrOfFactor" () args+ => ArgsHMap "lrOfFactor" () args -> FactorScheduler lrOfFactor args = Factor base factor step stop- where + where factor = args ! #factor step = args ! #step base = fromMaybe 0.01 (args !? #base)@@ -43,16 +43,16 @@ getLR (Multifactor base factor steps) nup = base * factor ^ (index nup steps) where index a bs = go a bs (0 :: Int)- go _ [] n = n+ go _ [] n = n go a (b:bs) n = if b > a then n else go a bs (n+1) -type instance ParameterList "lrOfMultifactor" =+type instance ParameterList "lrOfMultifactor" () = '[ '("factor", 'AttrReq Float), '("steps", 'AttrReq [Int]), '("base", 'AttrOpt Float)] -lrOfMultifactor :: Fullfilled "lrOfMultifactor" args- => ArgsHMap "lrOfMultifactor" args -> MultifactorScheduler+lrOfMultifactor :: Fullfilled "lrOfMultifactor" () args+ => ArgsHMap "lrOfMultifactor" () args -> MultifactorScheduler lrOfMultifactor args = Multifactor base factor steps- where + where factor = args ! #factor steps = args ! #steps base = fromMaybe 0.01 (args !? #base)@@ -60,17 +60,17 @@ data PolyScheduler = Poly Float Float Int instance LrScheduler PolyScheduler where getLR (Poly base power maxnup) nup =- if nup < maxnup + if nup < maxnup then base * (1 - fromIntegral nup / fromIntegral maxnup) ** power else 0 -type instance ParameterList "lrOfPoly" =+type instance ParameterList "lrOfPoly" () = '[ '("maxnup", 'AttrReq Int), '("power", 'AttrReq Float), '("base", 'AttrOpt Float)] -lrOfPoly :: Fullfilled "lrOfPoly" args- => ArgsHMap "lrOfPoly" args -> PolyScheduler+lrOfPoly :: Fullfilled "lrOfPoly" () args+ => ArgsHMap "lrOfPoly" () args -> PolyScheduler lrOfPoly args = Poly base power maxnup- where + where maxnup = args ! #maxnup base = fromMaybe 0.01 (args !? #base) power = fromMaybe 2 (args !? #power)
+ src/MXNet/NN/Module.hs view
@@ -0,0 +1,291 @@+module MXNet.NN.Module where++import Control.Lens (has, ix, use, (%=), (+=), (.=), (^?!))+import Formatting (int, sformat, stext, (%))+import GHC.TypeLits (KnownSymbol)+import RIO+import qualified RIO.HashMap as M+import qualified RIO.HashSet as S+import RIO.List (zipWith3)+import qualified RIO.NonEmpty as NE+import qualified RIO.Text as T+import qualified RIO.Vector.Storable as VS++import MXNet.Base+import MXNet.NN.DataIter.Class (Dataset (..), DatasetProp (..))+import MXNet.NN.EvalMetric (EvalMetricMethod (..), MetricData)+import MXNet.NN.Layer (copy)+import MXNet.NN.Optimizer (Optimizer, optimize)+import MXNet.NN.Session (withSession)+import MXNet.NN.TaggedState (Tagged (..), untag)+import MXNet.NN.Types+++data InitError = UnknownShape Text+ | ShouldNotBeScalar Text+ | ShapeNotAgree Text (NonEmpty Int) (NonEmpty Int)+ | BadShapeValue Text Text+ | BadInitValue Text Text+ deriving (Typeable, Show)+instance Exception InitError+++initialize :: forall tag dty. (HasCallStack, FloatDType dty) => SymbolHandle -> Config dty -> IO (TaggedModuleState dty tag)+initialize symbol config = do+ -- give a initial batch_size = 1 for the placeholders+ let spec1 = M.difference input_shapes initializers+ spec2 = initializers+ dinit = config ^. cfg_default_initializer+ cxt = config ^. cfg_context+ -- remove any input or label from the fixed set+ fixed = (config ^. cfg_fixed_params) `S.difference`+ (S.fromList $ M.keys input_shapes ++ label_names)++ -- although it is possible to set __shape__ to certain symbols, we+ -- don't use the information for shape inference right now.++ (args, _, auxs, _) <- inferShape symbol (M.toList spec1)+ arg_with_shp <- M.fromList <$> mapM checkTensorShape args+ aux_with_shp <- M.fromList <$> mapM checkTensorShape auxs+ let all_shapes = arg_with_shp `M.union` aux_with_shp++ attrs <- mxSymbolListAttr symbol+ let node_with_init = M.filter (has $ ix "__init__") attrs+ ---------------------+ -- important! labels should be merged into placeholders,+ -- otherwise the labels are considered to have gradient.+ ---------------------+ let lbl_with_shp = M.filterWithKey (\k _ -> k `elem` label_names) arg_with_shp+ phl_with_shp = M.map _shape_nonempty spec1 `M.union` lbl_with_shp+ placeholders <- mapM (flip makeEmptyNDArray cxt) phl_with_shp++ arg_init <- M.traverseWithKey (initN arg_with_shp aux_with_shp fixed) node_with_init+ arg_tensors <- M.traverseWithKey (initI placeholders fixed spec2 dinit) arg_with_shp+ aux_tensors <- M.traverseWithKey (initA dinit) aux_with_shp++ let params = arg_init `M.union` arg_tensors `M.union` aux_tensors+ --forM (M.toList params) $ \(k, p) -> do+ -- o <- case p of+ -- ParameterV a -> fmap (\s -> ("V", [s])) $ ndshape a+ -- ParameterF a -> fmap (\s -> ("F", [s])) $ ndshape a+ -- ParameterG a b -> liftM2 (\s t -> ("V", [s, t])) (ndshape a) (ndshape b)+ -- ParameterA a -> fmap (\s -> ("A", [s])) $ ndshape a+ -- traceShowM (" param", k, o)++ executor <- bind symbol cxt params True++ return $ Tagged $ ModuleState {+ _mod_symbol = symbol,+ _mod_input_shapes = phl_with_shp,+ _mod_params = params,+ _mod_context = cxt,+ _mod_executor = executor,+ _mod_statistics = Statistics 0 0,+ _mod_scores = M.empty,+ _mod_fixed_args = fixed+ }+ where+ input_shapes = config ^. cfg_data+ label_names = config ^. cfg_label+ initializers = config ^. cfg_initializers+ context = config ^. cfg_context+ -- initialize input symbols.+ -- placeholders are backed by empty NDArray,+ -- other input symbols are initialized by an initializer.+ initI placeholders fixed spec2 dinit inp shp = do+ case M.lookup inp placeholders of+ Just in_arg -> do+ return $ ParameterV in_arg+ Nothing -> do+ arg_in <- case M.lookup inp spec2 of+ Just cinit -> cinit inp shp context+ Nothing -> dinit inp shp context+ if S.member inp fixed+ then return $ ParameterF arg_in+ else do+ arg_gr <- makeEmptyNDArray shp context+ return $ ParameterG arg_in arg_gr+ -- initialize auxiliary symbols.+ initA dinit aux shp = do+ arg_aux <- dinit aux shp context+ return $ ParameterA arg_aux++ -- initialize symbol with __init__ attribute+ initN args auxs fixed name attrs = do+ let shp1 = M.lookup name (M.union args auxs)+ shp2 = M.lookup "__shape__" attrs++ shp2 <- case shp2 of+ Nothing -> pure Nothing+ Just s -> case readMaybe (T.unpack s) >>= NE.nonEmpty of+ Nothing -> throwM $ BadShapeValue name s+ t -> pure t++ shape <- case (shp1, shp2) of+ (Nothing, Nothing) -> throwM $ UnknownShape name+ (Just shp1, Just shp2) | shp1 /= shp2 -> throwM $ ShapeNotAgree name shp1 shp2+ (Just shp1, _) -> pure shp1+ (_, Just shp2) -> pure shp2++ let value_text = attrs ^. ix "__init__"+ array <- case T.unpack value_text of+ "[\"zero\", {}]" -> zeros shape >>= flip toContext context+ "[\"one\", {}]" -> ones shape >>= flip toContext context+ v -> case readMaybe v of+ Nothing -> throwM $ BadInitValue name value_text+ Just v -> makeNDArray shape context $ VS.fromList v++ case (S.member name fixed, M.member name args, M.member name auxs) of+ (True, _, _) -> return $ ParameterF array+ (False, False, True) -> return $ ParameterA array+ (False, True, False) -> do+ grad <- makeEmptyNDArray shape context+ return $ ParameterG array grad++ checkTensorShape (name, SScalar) = throwIO $ ShouldNotBeScalar name+ checkTensorShape (name, STensor s) = return (name, s)+++bind :: (HasCallStack, FloatDType dty) => SymbolHandle -> Context -> M.HashMap Text (Parameter dty) -> Bool -> IO (Executor dty)+bind symbol context params trainable = do+ argnames <- listArguments symbol+ auxnames <- listAuxiliaryStates symbol+ -- sanity check+ assert (S.fromList (M.keys params) == S.fromList (argnames ++ auxnames)) (return ())+ -- the parameters to bind should be arranged in the same order as the argnames+ let num_args = length argnames+ arg_all = map ((params ^?!) . ix) argnames+ arg_in = map (\case {+ ParameterV t -> t;+ ParameterF t -> t;+ ParameterG t _ -> t;+ ParameterA _ -> error "auxiliary parameter shouldn't occur"+ }) arg_all+ arg_gr_w_req = if trainable then map (\case {ParameterG _ t -> Just (t, 1); _ -> Nothing}) arg_all+ else replicate num_args Nothing+ aux_arg_aux = map (_param_aux . (params ^?!) . ix) auxnames+ execBind symbol context arg_in arg_gr_w_req aux_arg_aux+++adapt :: (FloatDType dty, MonadIO m) => M.HashMap Text (NDArray dty) -> Module tag dty m ()+adapt inputs = do+ symbol <- use $ untag . mod_symbol+ exec <- use $ untag . mod_executor+ context <- use $ untag . mod_context+ fixed <- use $ untag . mod_fixed_args+ shapes <- use $ untag . mod_input_shapes+ shapes' <- liftIO $ mapM ndshape inputs++ -- reshape the executor (and arg, aux arrays)+ when (shapes /= shapes') $ do+ (args, grads, auxs, exec) <- liftIO $ execReshapeEx exec True True context (M.toList shapes')+ arg_names <- liftIO $ listArguments symbol+ aux_names <- liftIO $ listAuxiliaryStates symbol+ let buildArg key a Nothing | S.member key fixed = (key, ParameterF a)+ buildArg key a Nothing | otherwise = (key, ParameterV a)+ buildArg key a (Just b) = (key, ParameterG a b)+ buildAux key a = (key, ParameterA a)+ arg_ndarrs = M.fromList $ zipWith3 buildArg arg_names args grads+ aux_ndarrs = M.fromList $ zipWith buildAux aux_names auxs++ -- it should be safe to free the old executor+ old_executor <- use (untag . mod_executor)+ liftIO $ execFree old_executor++ untag . mod_executor .= exec+ untag . mod_input_shapes .= shapes'+ untag . mod_params .= M.union arg_ndarrs aux_ndarrs++ -- copy the ndarray+ targets <- use $ untag . mod_params+ forM_ (M.toList inputs) $ \ (k, src) -> liftIO $ do+ case M.lookup k targets of+ Just (ParameterV dst) -> void $ copy src dst+ _ -> return ()++forwardOnly :: (FloatDType dty, MonadIO m) => M.HashMap Text (NDArray dty) -> Module tag dty m [NDArray dty]+forwardOnly inputs = do+ adapt inputs+ exec <- use $ untag . mod_executor+ liftIO $ do+ execForward exec False+ execGetOutputs exec++fit :: (FloatDType dty, MonadIO m) => M.HashMap Text (NDArray dty) -> Module tag dty m ()+fit inputs = do+ adapt inputs+ exec <- use $ untag . mod_executor+ liftIO $ do+ execForward exec True+ execBackward exec []++update :: (Optimizer opt, FloatDType dty, MonadIO m) => opt dty -> M.HashMap Text any -> Module tag dty m ()+update opt blacklist = do+ params <- use $ untag . mod_params+ forM_ (M.toList params) $ \case+ (k, ParameterG weig grad) | not (M.member k blacklist) -> do+ optimize opt k weig grad+ _ -> return ()+ untag . mod_statistics . stat_num_upd += 1+++fitAndEval :: (FloatDType dty, Optimizer opt, EvalMetricMethod mtr, MonadIO m)+ => opt dty -> M.HashMap Text (NDArray dty) -> MetricData mtr dty -> Module tag dty m ()+fitAndEval opt datAndLbl metric = do+ fit datAndLbl+ update opt M.empty+ exec <- use $ untag . mod_executor+ out <- liftIO $ execGetOutputs exec+ eval_results <- metricUpdate metric datAndLbl out+ untag . mod_scores %= M.union eval_results+++fitDataset :: (KnownSymbol tag, Dataset d, DatasetProp d e, FloatDType a,+ MonadIO m, MonadThrow m, MonadReader env m, HasLogFunc env,+ HasCallStack,+ DatasetMonadConstraint d m,+ Optimizer opt, EvalMetricMethod mtr)+ => TaggedModuleState a tag+ -> d m e+ -> d m e+ -> ([Text] -> e -> M.HashMap Text (NDArray a))+ -> opt a+ -> mtr a+ -> Int+ -> m ()+fitDataset sess trainDataset valDataset make_binding opt metric epochs = do+ -- callbacks <- use sess_callbacks++ sess_mvar <- newMVar sess+ let variables = M.keys $ sess ^. untag . mod_input_shapes+ total <- sizeD trainDataset+ -- batchSize <- batchSizeD trainDataset >>= maybe (throwM DatasetOfUnknownBatchSize) return++ logInfo $ display ("[Train]" :: Text)+ forM_ [1..epochs] $ \epochInd -> do+ trainMetricData <- newMetric "train" metric++ logInfo . display $ sformat ("epoch " % int) epochInd+ -- forM_ callbacks (begOfEpoch epochInd total)++ void $ forEachD_i trainDataset $ \(i, item) -> withSession sess_mvar $ do+ -- forM_ callbacks (begOfBatch i batchSize)+ let binding = make_binding variables item+ fitAndEval opt binding trainMetricData+ eval <- metricFormat trainMetricData+ logInfo . display $ sformat (int % int % stext) i total eval+ -- forM_ callbacks (endOfBatch i batchSize)++ -- forM_ callbacks (endOfEpoch epochInd total)++ logInfo $ display ("[Validate]" :: Text)+ valMetricData <- newMetric "val" metric+ void $ forEachD valDataset $ \item -> withSession sess_mvar $ do+ let binding = make_binding variables item+ -- TODO: it is bad to pass labels to forwardOnly+ out <- forwardOnly binding+ metricUpdate valMetricData binding out+ eval <- metricFormat valMetricData+ logInfo $ display eval+ --+ -- forM_ callbacks (endOfVal epochInd total)
− src/MXNet/NN/NDArray.hs
@@ -1,19 +0,0 @@-module MXNet.NN.NDArray where--import MXNet.Base-import qualified MXNet.Base.Operators.NDArray as I--reshape :: DType a => NDArray a -> [Int] -> IO (NDArray a)-reshape arr shp = do- [hdl] <- I._Reshape (#data := unNDArray arr .& #shape := shp .& Nil)- return $ NDArray hdl--transpose :: DType a => NDArray a -> [Int] -> IO (NDArray a)-transpose arr axes = do- [hdl] <- I.transpose (#data := unNDArray arr .& #axes := axes .& Nil)- return $ NDArray hdl--copy :: DType a => NDArray a -> NDArray a -> IO (NDArray a)-copy src dst = do- I._copyto_upd [unNDArray dst] (#data := unNDArray src .& Nil)- return dst
src/MXNet/NN/Optimizer.hs view
@@ -1,26 +1,22 @@-{-# LANGUAGE TypeOperators #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE OverloadedLists #-} {-# LANGUAGE UndecidableInstances #-}--module MXNet.NN.Optimizer (- Optimizer(..),- OptimizerTag(..)-) where+module MXNet.NN.Optimizer where -import MXNet.Base hiding (Symbol)-import qualified MXNet.Base.Operators.NDArray as A+import Control.Lens (use, (.=))+import GHC.Exts (Constraint)+import GHC.TypeLits+import RIO+import qualified RIO.HashMap as M+import RIO.State -import Data.IORef-import GHC.TypeLits-import GHC.Exts (Constraint)-import qualified Data.HashMap.Strict as M-import Control.Monad.IO.Class (MonadIO, liftIO)-import Control.Monad.State.Class (MonadState)-import Control.Lens (use, (.=))-import MXNet.NN.LrScheduler (LrScheduler(..))-import MXNet.NN.Types (Statistics, stat_num_upd, stat_last_lr)+import MXNet.Base hiding (Symbol)+import qualified MXNet.Base.Operators.Tensor as T+import MXNet.NN.LrScheduler (LrScheduler (..))+import MXNet.NN.TaggedState (untag)+import MXNet.NN.Types (TaggedModuleState, mod_statistics,+ stat_last_lr, stat_num_upd) -- | Abstract Optimizer type class class Optimizer (opt :: * -> *) where@@ -30,13 +26,15 @@ -- | Specific optional arguments -- type OptArgsList opt :: [KV *] -- | make the optimizer- makeOptimizer :: (DType dtype, LrScheduler sch, OptimizerCst opt dtype args) - => OptimizerTag opt -> sch -> ArgsHMap (OptimizerSym opt) args -> IO (opt dtype)+ makeOptimizer :: (DType dtype, LrScheduler sch, OptimizerCst opt dtype args, MonadIO m)+ => OptimizerTag opt -> sch+ -> ArgsHMap (OptimizerSym opt) (NDArray dtype) args+ -> m (opt dtype) -- | run the optimizer with the input & expected tensor- optimize :: (DType dtype, MonadIO m, MonadState Statistics m) + optimize :: (DType dtype, MonadState (TaggedModuleState dtype t) m, MonadIO m) => opt dtype -- optimizer- -> String -- symbol name to optimize- -> NDArray dytpe -- parameter+ -> Text -- symbol name to optimize+ -> NDArray dtype -- parameter -> NDArray dtype -- gradient -> m () @@ -46,32 +44,37 @@ -- | SGD optimizer data SGD_Opt dtype where SGD_Opt :: (LrScheduler sch, OptimizerCst SGD_Opt dtype args)- => sch -> ArgsHMap (OptimizerSym SGD_Opt) args -> SGD_Opt dtype+ => sch -> ArgsHMap (OptimizerSym SGD_Opt) (NDArray dtype) args+ -> SGD_Opt dtype -type instance OptimizerSym SGD_Opt = "sgd_update(ndarray)"+type instance OptimizerSym SGD_Opt = "_sgd_update" -- 1.0.0 type instance OptimizerCst SGD_Opt dt args = HasArgs (OptimizerSym SGD_Opt) args '["wd", "rescale_grad", "clip_gradient"]-type instance OptimizerCst SGD_Opt dt args = HasArgs (OptimizerSym SGD_Opt) args '["wd", "rescale_grad", "clip_gradient", "lazy_update"]+type instance OptimizerCst SGD_Opt dt args =+ HasArgs (OptimizerSym SGD_Opt) (NDArray dt) args '["wd", "rescale_grad", "clip_gradient", "lazy_update"] instance Optimizer SGD_Opt where data OptimizerTag SGD_Opt = SGD makeOptimizer SGD sch args = return $ SGD_Opt sch args- optimize (SGD_Opt sch args) _ (NDArray weight) (NDArray gradient) = do- nup <- use stat_num_upd+ optimize (SGD_Opt sch args) _ weight gradient = do+ nup <- use $ untag . mod_statistics . stat_num_upd let lr = getLR sch nup- stat_last_lr .= lr- liftIO $ A.sgd_update_upd [weight] (- #weight := weight .& - #grad := gradient .& - #lr := lr .& args)+ untag . mod_statistics . stat_last_lr .= lr+ liftIO $ void $ T._sgd_update (#weight := weight+ .& #grad := gradient+ .& #lr := lr .& args)+ (Just [weight]) -- | SGD with momentum optimizer data SGD_Mom_Opt dtype where SGD_Mom_Opt :: (LrScheduler sch, OptimizerCst SGD_Mom_Opt dtype args)- => sch -> ArgsHMap (OptimizerSym SGD_Mom_Opt) args -> (IORef (M.HashMap String (NDArray dtype))) -> SGD_Mom_Opt dtype+ => sch -> ArgsHMap (OptimizerSym SGD_Mom_Opt) (NDArray dtype) args+ -> (IORef (M.HashMap Text (NDArray dtype)))+ -> SGD_Mom_Opt dtype -type instance OptimizerSym SGD_Mom_Opt = "sgd_mom_update(ndarray)"+type instance OptimizerSym SGD_Mom_Opt = "_sgd_mom_update" -- 1.0.0 type instance OptimizerCst SGD_Mom_Opt dt args = HasArgs (OptimizerSym SGD_Mom_Opt) args '["momentum", "wd", "rescale_grad", "clip_gradient"]-type instance OptimizerCst SGD_Mom_Opt dt args = HasArgs (OptimizerSym SGD_Mom_Opt) args '["momentum", "wd", "rescale_grad", "clip_gradient", "lazy_update"]+type instance OptimizerCst SGD_Mom_Opt dt args =+ HasArgs (OptimizerSym SGD_Mom_Opt) (NDArray dt) args '["momentum", "wd", "rescale_grad", "clip_gradient", "lazy_update"] instance Optimizer SGD_Mom_Opt where data OptimizerTag SGD_Mom_Opt = SGD'Mom@@ -79,32 +82,42 @@ empty <- newIORef M.empty return $ SGD_Mom_Opt sch args empty - optimize (SGD_Mom_Opt sch args emaref) symbol (NDArray weight) (NDArray gradient) = do- nup <- use stat_num_upd+ optimize (SGD_Mom_Opt sch args emaref) symbol weight gradient = do+ nup <- use $ untag . mod_statistics . stat_num_upd let lr = getLR sch nup- stat_last_lr .= lr+ untag . mod_statistics . stat_last_lr .= lr liftIO $ do ema <- readIORef emaref momentum <- case M.lookup symbol ema of Nothing -> do- [mom] <- A.zeros_like (#data := weight .& Nil)- writeIORef emaref (M.insert symbol (NDArray mom) ema)+ mom <- prim T._zeros_like (#data := weight .& Nil)+ writeIORef emaref (M.insert symbol mom ema) return mom- Just (NDArray a) -> return a- A.sgd_mom_update_upd [weight] (- #weight := weight .& - #grad := gradient .& - #mom := momentum .& - #lr := lr .& args)+ Just a -> return a+ -- let norm x = prim T._norm (#data := x .& #ord := 1 .& Nil)+ -- [w0] <- toVector =<< norm weight+ -- [m0] <- toVector =<< norm momentum+ -- [g0] <- toVector =<< norm gradient+ void $ T._sgd_mom_update (#weight := weight+ .& #grad := gradient+ .& #mom := momentum+ .& #lr := lr .& args)+ (Just [weight])+ -- [w1] <- toVector =<< norm weight+ -- [m1] <- toVector =<< norm momentum+ -- traceShowM ("opt", symbol, w0, m0, g0, w1, m1) -- | ADAM optmizer data ADAM_Opt dtype where- ADAM_Opt :: (LrScheduler sch, OptimizerCst ADAM_Opt dtype args) - => sch -> ArgsHMap (OptimizerSym ADAM_Opt) args -> IORef (M.HashMap String (NDArray dtype, NDArray dtype)) -> ADAM_Opt dtype+ ADAM_Opt :: (LrScheduler sch, OptimizerCst ADAM_Opt dtype args)+ => sch -> ArgsHMap (OptimizerSym ADAM_Opt) (NDArray dtype) args+ -> IORef (M.HashMap Text (NDArray dtype, NDArray dtype))+ -> ADAM_Opt dtype -type instance OptimizerSym ADAM_Opt = "adam_update(ndarray)"+type instance OptimizerSym ADAM_Opt = "_adam_update" -- 1.0.0 type instance OptimizerCst ADAM_Opt dt args = HasArgs (OptimizerSym ADAM_Opt) args '["beta1", "beta2", "epsilon", "wd", "rescale_grad", "clip_gradient"]-type instance OptimizerCst ADAM_Opt dt args = HasArgs (OptimizerSym ADAM_Opt) args '["beta1", "beta2", "epsilon", "wd", "rescale_grad", "clip_gradient", "lazy_update"]+type instance OptimizerCst ADAM_Opt dt args =+ HasArgs (OptimizerSym ADAM_Opt) (NDArray dt) args '["beta1", "beta2", "epsilon", "wd", "rescale_grad", "clip_gradient", "lazy_update"] instance Optimizer ADAM_Opt where data OptimizerTag ADAM_Opt = ADAM@@ -112,22 +125,66 @@ empty <- newIORef M.empty return $ ADAM_Opt sch args empty - optimize (ADAM_Opt sch args emaref) symbol (NDArray weight) (NDArray gradient) = do- nup <- use stat_num_upd+ optimize (ADAM_Opt sch args emaref) symbol weight gradient = do+ nup <- use $ untag . mod_statistics . stat_num_upd let lr = getLR sch nup- stat_last_lr .= lr+ untag . mod_statistics . stat_last_lr .= lr liftIO $ do ema <- readIORef emaref (moving_avg, moving_var) <- case M.lookup symbol ema of Nothing -> do- [avg] <- A.zeros_like (#data := weight .& Nil)- [var] <- A.zeros_like (#data := weight .& Nil)- writeIORef emaref (M.insert symbol (NDArray avg, NDArray var) ema)+ avg <- prim T._zeros_like (#data := weight .& Nil)+ var <- prim T._ones_like (#data := weight .& Nil)+ writeIORef emaref (M.insert symbol (avg, var) ema) return (avg, var)- Just (NDArray a, NDArray v) -> return (a, v)- A.adam_update_upd [weight] (- #weight := weight .&- #grad := gradient .&- #mean := moving_avg .&- #var := moving_var .&- #lr := lr .& args)+ Just (a, v) -> return (a, v)+ void $ T._adam_update (#weight := weight+ .& #grad := gradient+ .& #mean := moving_avg+ .& #var := moving_var+ .& #lr := lr .& args)+ (Just [weight])+++#ifdef MXNET_GEQ_10600++data ADAMW_Opt dtype where+ ADAMW_Opt :: (LrScheduler sch, OptimizerCst ADAMW_Opt dtype args)+ => sch -> ArgsHMap (OptimizerSym ADAMW_Opt) (NDArray dtype) args+ -> IORef (M.HashMap Text (NDArray dtype, NDArray dtype))+ -> ADAMW_Opt dtype++type instance OptimizerSym ADAMW_Opt = "__adamw_update"+type instance OptimizerCst ADAMW_Opt dt args =+ HasArgs (OptimizerSym ADAMW_Opt) (NDArray dt) args+ '["beta1", "beta2", "epsilon", "wd", "eta", "clip_gradient", "rescale_grad"]++instance Optimizer ADAMW_Opt where+ data OptimizerTag ADAMW_Opt = ADAMW+ makeOptimizer ADAMW sch args = do+ empty <- newIORef M.empty+ return $ ADAMW_Opt sch args empty++ optimize (ADAMW_Opt sch args emaref) symbol weight gradient = do+ nup <- use $ untag . mod_statistics . stat_num_upd+ let lr = getLR sch nup+ untag . mod_statistics . stat_last_lr .= lr+ liftIO $ do+ ema <- readIORef emaref+ (moving_avg, moving_var) <- case M.lookup symbol ema of+ Nothing -> do+ avg <- prim T._zeros_like (#data := weight .& Nil)+ var <- prim T._ones_like (#data := weight .& Nil)+ writeIORef emaref (M.insert symbol (avg, var) ema)+ return (avg, var)+ Just (a, v) -> return (a, v)+ void $ T.__adamw_update+ (#weight := weight .&+ #grad := gradient .&+ #mean := moving_avg .&+ #var := moving_var .&+ #lr := lr .& args)+ (Just [weight])++#endif+
+ src/MXNet/NN/Session.hs view
@@ -0,0 +1,115 @@+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE FunctionalDependencies #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE TemplateHaskell #-}+module MXNet.NN.Session where++import Data.Kind (Constraint)+import qualified Data.Type.Index as DT+import qualified Data.Type.Product as DT+import qualified GHC.TypeLits as L+import RIO+import qualified RIO.State as ST++import MXNet.Base+import MXNet.NN.TaggedState (liftSub, toPair)+import MXNet.NN.Types (Module, ModuleSet, ModuleState,+ TaggedModuleState)++class Session (sess :: (* -> *) -> * -> *) (sst :: *) | sess -> sst, sst -> sess where+ type SessionDType sst+ type SessionHasModule sst (t :: L.Symbol) :: Constraint+ runSession :: sess m r -> sst -> m (r, sst)+ subSession :: MonadIO m => SessionHasModule sst t => Module t (SessionDType sst) m r -> sess m r+ sessionStates :: MonadIO m => sess m [(String, ModuleState (SessionDType sst))]++class HasSessionRef e s | e -> s where+ sessionRefL :: Lens' e (MVar s)++askSession :: (MonadIO m, MonadReader e m, HasSessionRef e s, Session sess s)+ => sess m r -> m r+askSession proc = do+ env <- ask+ let var = env ^. sessionRefL+ st <- liftIO $ takeMVar var+ (rt, st) <- runSession proc st+ liftIO $ putMVar var st+ return rt++withSession :: (MonadIO m, Session sess s)+ => MVar s -> sess m r -> m r+withSession var proc = do+ st <- liftIO $ takeMVar var+ (rt, st) <- runSession proc st+ liftIO $ putMVar var st+ return rt++instance (DT.Every L.KnownSymbol t) => Session (ModuleSet t a) (DT.Prod (TaggedModuleState a) t) where+ type SessionDType (DT.Prod (TaggedModuleState a) t) = a+ type SessionHasModule (DT.Prod (TaggedModuleState a) t) t' = DT.Elem t t'+ runSession = ST.runStateT+ subSession = liftSub+ sessionStates = walk <$> ST.get+ where+ walk :: DT.Every L.KnownSymbol t => DT.Prod (TaggedModuleState a) t -> [(String, ModuleState a)]+ walk DT.Ø = []+ walk (v DT.:< rest) = toPair v : walk rest++instance (L.KnownSymbol t) => Session (Module t a) (TaggedModuleState a t) where+ type SessionDType (TaggedModuleState a t) = a+ type SessionHasModule (TaggedModuleState a t) t' = t ~ t'+ runSession = ST.runStateT+ subSession = id+ sessionStates = (:[]) . toPair <$> ST.get++class CallbackClass c where+ begOfBatch :: ( L.KnownSymbol t+ , DType a+ , MonadIO m+ , MonadReader env m+ , HasLogFunc env+ , HasCallStack)+ => Int -> Int -> c -> Module t a m ()+ begOfBatch _ _ _ = return ()+ endOfBatch :: ( L.KnownSymbol t+ , DType a+ , MonadIO m+ , MonadReader env m+ , HasLogFunc env+ , HasCallStack)+ => Int -> Int -> c -> Module t a m ()+ endOfBatch _ _ _ = return ()+ begOfEpoch :: ( L.KnownSymbol t+ , DType a+ , MonadIO m+ , MonadReader env m+ , HasLogFunc env+ , HasCallStack)+ => Int -> Int -> c -> Module t a m ()+ begOfEpoch _ _ _ = return ()+ endOfEpoch :: ( L.KnownSymbol t+ , DType a+ , MonadIO m+ , MonadReader env m+ , HasLogFunc env+ , HasCallStack)+ => Int -> Int -> c -> Module t a m ()+ endOfEpoch _ _ _ = return ()+ endOfVal :: ( L.KnownSymbol t+ , DType a+ , MonadIO m+ , MonadReader env m+ , HasLogFunc env+ , HasCallStack)+ => Int -> Int -> c -> Module t a m ()+ endOfVal _ _ _ = return ()+data Callback where+ Callback :: CallbackClass a => a -> Callback++instance CallbackClass Callback where+ begOfBatch i n (Callback a) = begOfBatch i n a+ endOfBatch i n (Callback a) = endOfBatch i n a+ begOfEpoch i n (Callback a) = begOfEpoch i n a+ endOfEpoch i n (Callback a) = endOfEpoch i n a+ endOfVal i n (Callback a) = endOfVal i n a+
+ src/MXNet/NN/TaggedState.hs view
@@ -0,0 +1,50 @@+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE DataKinds, TypeOperators #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE PolyKinds #-}+module MXNet.NN.TaggedState where++import RIO+import RIO.State (StateT(..))+import qualified GHC.TypeLits as L+import Data.Type.Product+import Data.Type.Index+import Control.Lens (makeLenses)+import Data.Proxy (Proxy(..))++newtype Tagged a (t :: L.Symbol) = Tagged {_untag :: a} deriving Show+makeLenses ''Tagged+++-- liftSub :: forall k (f :: k -> *) ss t m a. (Elem ss t, Monad m) => ReaderT (f t) m a -> ReaderT (Prod f ss) m a+-- liftSub (ReaderT m) = ReaderT (m . index elemIndex)+liftSub :: forall k (f :: k -> *) s1 s2 m a. (Elem s2 s1, Monad m) => StateT (f s1) m a -> StateT (Prod f s2) m a+liftSub (StateT m1) = StateT $ \s -> do+ (a, si) <- m1 $ index elemIndex s+ let new_s = modify elemIndex si s+ new_s `seq` return (a, new_s)+++modify :: Index as a -> f a -> Prod f as -> Prod f as+modify IZ new (_ :< remainder) = new :< remainder+modify (IS s) new (item :< remainder) = item :< modify s new remainder+++toPair :: forall t a. L.KnownSymbol t => Tagged a t -> (String, a)+toPair (Tagged a)= (L.symbolVal (Proxy :: Proxy t), a)+++-- a1 :: StateT (Tagged Int "A") IO ()+-- a1 = put (Tagged 4)+--+-- a2 :: StateT (Tagged String "B") IO ()+-- a2 = put (Tagged "hi")+--+-- a3 :: StateT (ProdI '[Tagged Int "A", Tagged String "B"]) IO ()+-- a3 = do+-- liftT a1+-- liftT a2++-- runStateT a3 (Identity (Tagged 0) :> Identity (Tagged ""))+
src/MXNet/NN/Types.hs view
@@ -1,99 +1,94 @@-{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# 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 Control.Exception.Base (Exception) -import Data.Typeable (Typeable) -import Data.Dynamic (Dynamic) -import Control.Monad.IO.Class (MonadIO) - -import MXNet.Base - --- | A parameter is two 'NDArray' to back a 'Symbol' -data Parameter a = ParameterI { _param_in :: NDArray a, _param_grad :: Maybe (NDArray a) } - | ParameterA { _param_aux :: NDArray a } - -- deriving Show - -data Statistics = Statistics { - _stat_num_upd :: !Int, - _stat_last_lr :: !Float -} - -class CallbackClass a where - begOfBatch :: (MonadIO m, DType e) => Int -> Int -> a -> TrainM e m () - begOfBatch _ _ _ = return () - endOfBatch :: (MonadIO m, DType e) => Int -> Int -> a -> TrainM e m () - endOfBatch _ _ _ = return () - begOfEpoch :: (MonadIO m, DType e) => Int -> Int -> a -> TrainM e m () - begOfEpoch _ _ _ = return () - endOfEpoch :: (MonadIO m, DType e) => Int -> Int -> a -> TrainM e m () - endOfEpoch _ _ _ = return () - endOfVal :: (MonadIO m, DType e) => Int -> Int -> a -> TrainM e m () - endOfVal _ _ _ = return () -data Callback where - Callback :: CallbackClass a => a -> Callback - -instance CallbackClass Callback where - begOfBatch i n (Callback a) = begOfBatch i n a - endOfBatch i n (Callback a) = endOfBatch i n a - begOfEpoch i n (Callback a) = begOfEpoch i n a - endOfEpoch i n (Callback a) = endOfEpoch i n a - endOfVal i n (Callback a) = endOfVal i n a +import Control.Lens (makeLenses) +import qualified Data.Type.Product as DT +import Data.Typeable (Typeable) +import RIO +import RIO.HashMap (HashMap) +import RIO.HashSet (HashSet) +import RIO.State (StateT) --- | Session is all the 'Parameters' and a 'Device' --- type Session a = (M.HashMap String (Parameter a), Context) -data Session a = Session { - _sess_symbol :: Symbol a - , _sess_data :: M.HashMap String [Int] - , _sess_label :: [String] - , _sess_param :: !(M.HashMap String (Parameter a)) - , _sess_context :: !Context - , _sess_callbacks :: [Callback] - , _sess_store :: M.HashMap String Dynamic - -- , _sess_prof :: (NominalDiffTime, NominalDiffTime, NominalDiffTime, NominalDiffTime, NominalDiffTime, NominalDiffTime) -} --- | TrainM is a 'StateT' monad -type TrainM a m = ST.StateT (Session a) (ST.StateT Statistics m) +import MXNet.Base +import MXNet.NN.TaggedState (Tagged) -- | For every symbol in the neural network, it can be placeholder or a variable. --- therefore, a Config is to specify the shape of the placeholder and the +-- therefore, a Config is to specify the shape of the placeholder and the -- method to initialize the variables. --- --- Note that it is not right to specify a symbol as both placeholder and +-- +-- Note that it is not right to specify a symbol as both placeholder and -- initializer, although it is tolerated and such a symbol is considered -- as a variable. --- --- Note that any symbol not specified will be initialized with the +-- +-- Note that any symbol not specified will be initialized with the -- _cfg_default_initializer. -data Config a = Config { - _cfg_data :: M.HashMap String [Int], - _cfg_label :: [String], - _cfg_initializers :: M.HashMap String (Initializer a), - _cfg_default_initializer :: Initializer a, - _cfg_context :: Context -} +data Config a = Config + { _cfg_data :: HashMap Text FShape + , _cfg_label :: [Text] + , _cfg_initializers :: HashMap Text (Initializer a) + , _cfg_default_initializer :: Initializer a + , _cfg_fixed_params :: HashSet Text + , _cfg_context :: Context + } --- | Initializer is about how to create a NDArray from the symbol name and the given shape. --- --- Usually, it can be a wrapper of MXNet operators, such as @random_uniform@, @random_normal@, +-- | Initializer is about how to create a NDArray from the symbol name and the given shape. +-- +-- Usually, it can be a wrapper of MXNet operators, such as @random_uniform@, @random_normal@, -- @random_gamma@, etc.. -type Initializer a = String -> [Int] -> Context -> IO (NDArray a) +type Initializer a = Text -> NonEmpty Int -> Context -> IO (NDArray a) -- | Possible exception in 'TrainM' -data Exc = MismatchedShapeOfSym String [Int] [Int] - | MismatchedShapeInEval [Int] [Int] - | NotAParameter String - | InvalidArgument String - | InferredShapeInComplete - | DatasetOfUnknownBatchSize - | LoadSessionInvalidTensorName String - | LoadSessionMismatchedTensorKind String +data Exc = MismatchedShapeOfSym Text (NonEmpty Int) (NonEmpty Int) + | MismatchedShapeInEval (NonEmpty Int) (NonEmpty Int) + | NotAParameter Text + | InvalidArgument Text + | InferredShapeInComplete + | DatasetOfUnknownBatchSize + | LoadSessionInvalidTensorName Text + | LoadSessionMismatchedTensorKind Text deriving (Show, Typeable) instance Exception Exc -makeLenses ''Config +type TaggedModuleState dty = Tagged (ModuleState dty) +type Module tag dty = StateT (TaggedModuleState dty tag) +type ModuleSet tags dty = StateT (DT.Prod (TaggedModuleState dty) tags) + +data ModuleState a = ModuleState + { _mod_symbol :: SymbolHandle + , _mod_input_shapes :: HashMap Text (NonEmpty Int) + , _mod_params :: HashMap Text (Parameter a) + , _mod_context :: Context + , _mod_executor :: Executor a + , _mod_statistics :: Statistics + , _mod_scores :: HashMap Text Double + , _mod_fixed_args :: HashSet Text + } + +-- | A parameter is two 'NDArray' to back a 'Symbol' +data Parameter a = ParameterV + { _param_var :: NDArray a + } + | ParameterF + { _param_arg :: NDArray a + } + | ParameterG + { _param_arg :: NDArray a + , _param_grad :: NDArray a + } + | ParameterA + { _param_aux :: NDArray a + } + deriving Show + +data Statistics = Statistics + { _stat_num_upd :: !Int + , _stat_last_lr :: !Float + } + makeLenses ''Statistics -makeLenses ''Session+makeLenses ''ModuleState +makeLenses ''Config +
src/MXNet/NN/Utils.hs view
@@ -1,65 +1,136 @@-{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE PartialTypeSignatures #-}+{-# LANGUAGE RecordWildCards #-} module MXNet.NN.Utils where -import Data.List (intersperse)-import qualified Data.Text as T-import qualified Data.HashMap.Strict as M-import Control.Lens (use)-import Control.Monad (forM_)-import Control.Monad.IO.Class (MonadIO, liftIO)-import Control.Monad.Trans.Resource (MonadThrow(..))-import Text.Printf+import Control.Lens (use)+import Formatting+import RIO+import RIO.Directory (getModificationTime, listDirectory)+import RIO.FilePath (dropExtension, (</>))+import qualified RIO.HashMap as M+import RIO.List (lastMaybe, sortOn)+import qualified RIO.NonEmpty as RNE+import qualified RIO.Text as T -import MXNet.Base (- Context(..), DType, NDArray(..), Symbol(..), - HMap(..), (.&), ArgOf(..),- mxSymbolSaveToFile, mxNDArraySave, mxNDArrayLoad)-import MXNet.NN.Types-import qualified MXNet.Base.Operators.NDArray as A+import MXNet.Base (Context (..), DType, NDArray (..),+ ndshape)+import MXNet.Base.Raw (mxNDArrayLoad, mxNDArraySave,+ mxSymbolSaveToFile)+import MXNet.NN.Layer (copy)+import MXNet.NN.TaggedState (untag)+import MXNet.NN.Types (Module, Parameter (..), mod_params,+ mod_symbol) -- | format a shape-formatShape :: [Int] -> String-formatShape shape = concat $ ["("] ++ intersperse "," (map show shape) ++ [")"]+formatShape :: NonEmpty Int -> Text+formatShape shape = let sshape = RNE.intersperse "," (RNE.map tshow shape)+ in T.concat $ ["("] ++ RNE.toList sshape ++ [")"] -- | format a context-formatContext :: Context -> String-formatContext Context{..} = getDeviceName _device_type ++ "(" ++ show _device_id ++ ")"- where +formatContext :: Context -> Text+formatContext Context{..} = sformat (stext % "(" % int % ")") (getDeviceName _device_type) _device_id+ where+ getDeviceName :: Int -> Text getDeviceName 1 = "cpu" getDeviceName 2 = "gpu" getDeviceName 3 = "cpu_pinned" getDeviceName _ = error "formatContext: unknown device type" -endsWith :: String -> String -> Bool-endsWith s1 s2 = T.isSuffixOf (T.pack s1) (T.pack s2)+-- class Session s where+-- saveSession :: (String -> String) -> Bool -> ST.StateT s IO ()+--+-- instance L.KnownSymbol tag => Session (TaggedModuleState a tag) where+-- saveSession make_filename save_symbol = do+-- st <- ST.get+-- let name = L.symbolVal (Proxy :: Proxy tag)+-- liftIO $ saveState save_symbol (make_filename name) (st ^. untag)+--+-- instance DT.Every L.KnownSymbol tags => Session (DT.Prod (TaggedModuleState a) tags) where+-- saveSession make_filename save_symbol = do+-- tagged_states <- ST.get+-- let names = toNames (DT.map1 (const Proxy) tagged_states)+-- states = DT.toList (^. untag) tagged_states+-- liftIO $ zipWithM_ (saveState save_symbol) names states+-- where+-- toNames :: forall (t :: [L.Symbol]). DT.Every L.KnownSymbol t => DT.Prod Proxy t -> [String]+-- toNames DT.Ø = []+-- toNames (v DT.:< rem) = L.symbolVal v : toNames rem -saveSession :: (MonadIO m, DType a) => String -> TrainM a m ()-saveSession filename = do- dat_vars <- M.keys <$> use sess_data- lbl_vars <- use sess_label- params <- use sess_param- net <- use sess_symbol- let all_vars = dat_vars ++ lbl_vars- modelParams = map getModelParam $ M.toList $ M.filterWithKey (\k _ -> not (k `elem` all_vars)) params++saveState :: MonadIO m => Bool -> String -> Module t a m ()+saveState save_symbol name = do+ params <- use (untag . mod_params)+ symbol <- use (untag . mod_symbol)+ let modelParams = concatMap getModelParam $ M.toList params liftIO $ do- mxSymbolSaveToFile (filename ++ ".json") (unSymbol net)- mxNDArraySave (filename ++ ".params") modelParams+ when save_symbol $ mxSymbolSaveToFile (T.pack $ name ++ ".json") symbol+ mxNDArraySave (T.pack $ name ++ ".params") modelParams where- getModelParam (key, ParameterI a _) = ("arg:" ++ key, unNDArray a)- getModelParam (key, ParameterA a) = ("aux:" ++ key, unNDArray a)+ getModelParam (_, ParameterV _) = []+ getModelParam (key, ParameterF a) = [(key, unNDArray a)]+ getModelParam (key, ParameterA a) = [(key, unNDArray a)]+ getModelParam (key, ParameterG a g) =+ [(key, unNDArray a), (key `T.append` "__grad", unNDArray g)] -loadSession :: (MonadThrow m, MonadIO m, DType a) => String -> [String] -> TrainM a m ()-loadSession filename ignores = do- arrays <- liftIO $ mxNDArrayLoad (filename ++ ".params")- params <- use sess_param- forM_ arrays $ \(name, hdl) -> - case break (==':') name of- (_, "") -> throwM (LoadSessionInvalidTensorName name)- ("", _) -> throwM (LoadSessionInvalidTensorName name)- (typ, ':':key) -> - case (key `elem` ignores, typ, M.lookup key params) of- (True, _, _) -> return ()- (_, _, Nothing) -> liftIO $ putStrLn $ printf "Tensor %s is missing." name- (_, "arg", Just (ParameterI target grad)) -> liftIO $ A._copyto_upd [unNDArray target] (#data := hdl .& Nil)- (_, "aux", Just (ParameterA target)) -> liftIO $ A._copyto_upd [unNDArray target] (#data := hdl .& Nil)- _ -> throwM (LoadSessionMismatchedTensorKind name)+loadState :: (DType a, MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack)+ => String -> [Text] -> Module t a m ()+loadState weights_filename ignores = do+ arrays <- liftIO $ mxNDArrayLoad (T.pack $ weights_filename ++ ".params")+ params <- use (untag . mod_params)+ forM_ arrays $ \(name, hdl) -> do+ let nameIfGrad = T.stripSuffix "__grad" name+ nameIngore = name `elem` ignores+ param1 = M.lookup name params+ param2 = nameIfGrad >>= flip M.lookup params+ case (nameIngore, param1, nameIfGrad, param2) of+ (True, _, _, _) ->+ return ()+ (_, Nothing, _, Nothing) ->+ lift $ logInfo $ display $ sformat+ ("Tensor " % stext % " is ignored as missing in the model.") name+ (_, Nothing, Just name', Just (ParameterG _ target)) -> do+ checkShape name' (NDArray hdl) target+ liftIO $ void $ copy hdl (unNDArray target)+ (_, Nothing, Just name', Just _) -> do+ -- we silently ignore any missing grad,+ -- for it is too common if we load the model for inference+ return ()+ (_, Just (ParameterV _), _, _) ->+ logWarn . display $ sformat+ ("Tensor " % stext % " is ignored as a variable in the model.") name+ (_, Just (ParameterG target _), _, _) -> do+ checkShape name (NDArray hdl) target+ liftIO $ void $ copy hdl (unNDArray target)+ (_, Just (ParameterF target), _, _) -> do+ checkShape name (NDArray hdl) target+ liftIO $ void $ copy hdl (unNDArray target)+ (_, Just (ParameterA target), _, _) -> do+ checkShape name (NDArray hdl) target+ liftIO $ void $ copy hdl (unNDArray target)+ where+ checkShape :: (MonadReader env m, HasLogFunc env, MonadIO m, DType a)+ => Text -> NDArray a -> NDArray a -> m ()+ checkShape name arr1 arr2 = do+ shp1 <- liftIO $ ndshape $ arr1+ shp2 <- liftIO $ ndshape $ arr2+ when (shp1 /= shp2) $+ logWarn . display $ sformat+ ("variable (" % stext %+ ") has shape " % stext %+ ", different from that in saved state " % stext %+ ".") name (tshow shp2) (tshow shp1)++lastSavedState :: MonadIO m => Text -> Text -> m (Maybe FilePath)+lastSavedState dir prefix = liftIO $ do+ let sdir = T.unpack dir+ files <- listDirectory sdir+ let match name = T.isSuffixOf ".params" name && T.isPrefixOf prefix name+ param_files = filter (match . T.pack) files+ if null param_files+ then return Nothing+ else do+ mod_time <- mapM (getModificationTime . (sdir </>)) param_files+ case lastMaybe $ sortOn snd (zip param_files mod_time) of+ Nothing -> return Nothing+ Just (latest, _) -> return $ Just $ sdir </> dropExtension latest+
src/MXNet/NN/Utils/GraphViz.hs view
@@ -1,47 +1,42 @@-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE TypeApplications, TypeOperators, DataKinds #-}-{-# LANGUAGE QuasiQuotes #-}-{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE RecordWildCards #-} module MXNet.NN.Utils.GraphViz ( dotPlot,- dotGraph, + dotGraph, GV.GraphvizOutput(..) ) where -import Data.Aeson-import Data.Aeson.Types-import Data.ByteString.Lazy.Char8 (pack)-import qualified Data.Map as M-import Control.Exception.Base (Exception)-import Control.Monad.Catch(MonadThrow(..))-import Data.Typeable (Typeable)-import Data.Maybe-import Numeric (readHex)-import Text.Printf (printf)-import Control.Monad (forM_, when)-import qualified Data.Text.Lazy as T-import qualified Data.GraphViz as GV+import Data.Aeson+import Data.Aeson.Types+import qualified Data.GraphViz as GV import qualified Data.GraphViz.Attributes.Complete as GV-import qualified Data.GraphViz.Types.Monadic as GVM-import qualified Data.GraphViz.Types.Generalised as GVM+import qualified Data.GraphViz.Types.Generalised as GVM+import qualified Data.GraphViz.Types.Monadic as GVM+import Data.Typeable (Typeable)+import Formatting+import Numeric (readHex)+import RIO+import qualified RIO.Map as M+import RIO.Partial (fromJust)+import qualified RIO.Text as T+import qualified RIO.Text.Lazy as TL -import MXNet.Base+import MXNet.Base -- The program `dot` must be found in the PATH. -dotPlot :: DType a => Symbol a -> GV.GraphvizOutput -> FilePath -> IO ()+dotPlot :: SymbolHandle -> GV.GraphvizOutput -> FilePath -> IO () dotPlot sym output filepath = do gr <- dotGraph sym _ <- GV.addExtension (GV.runGraphvizCommand GV.Dot gr) output filepath return () -data JSNode = JSNode {- _node_op :: String,- _node_name :: String,- _node_attrs :: Maybe (M.Map String String),- _node_inputs :: [[Int]]-} deriving (Show)+data JSNode = JSNode+ { _node_op :: Text+ , _node_name :: Text+ , _node_attrs :: Maybe (M.Map Text Text)+ , _node_inputs :: [[Int]]+ }+ deriving (Show) instance FromJSON JSNode where parseJSON (Object v) = JSNode <$> v .: "op"@@ -50,9 +45,10 @@ <*> v .: "inputs" parseJSON invalid = typeMismatch "JSNode" invalid -data JSGraph = JSGraph {- _symbol_nodes :: [JSNode]-} deriving (Show)+data JSGraph = JSGraph+ { _symbol_nodes :: [JSNode]+ }+ deriving (Show) instance FromJSON JSGraph where parseJSON (Object v) = JSGraph <$> v .: "nodes"@@ -60,109 +56,110 @@ -- plot_network -- https://github.com/apache/incubator-mxnet/blob/master/python/mxnet/visualization.py#L196-dotGraph :: DType a => Symbol a -> IO (GVM.DotGraph Int)-dotGraph (Symbol sym) = do+dotGraph :: SymbolHandle -> IO (GVM.DotGraph Int)+dotGraph sym = do js <- mxSymbolSaveToJSON sym auxnodes <- mxSymbolListAuxiliaryStates sym- case eitherDecode $ pack js of+ case eitherDecodeStrict $ T.encodeUtf8 js of Left _ -> throwM CannotDecodeJSONofSymbol Right (JSGraph nodes) -> return $ GVM.digraph (GV.Num $ GV.Int 0) $ do let nodesWithIdx = (zip [0..] nodes)- blacklist = map fst $ - filter (\(_, node) -> elem (_node_name node) auxnodes || - _like "-weight" node || _like "-bias" node || - _like "-beta" node || _like "-gamma" node) + blacklist = map fst $+ filter (\(_, node) -> elem (_node_name node) auxnodes ||+ _like "-weight" node || _like "-bias" node ||+ _like "-beta" node || _like "-gamma" node) nodesWithIdx forM_ nodesWithIdx (mkNode_ blacklist) forM_ nodesWithIdx (mkEdge_ blacklist) where- mkNode_ blacklist (nodeid, JSNode{..}) = case _node_op of - "null" -> - when (not $ elem nodeid blacklist) $ - mkNode nodeid (#label := _node_name .& #shape := GV.Ellipse .& #fillcolor := colors !! 0 .& Nil)+ mkNode_ blacklist (nodeid, JSNode{..}) = case _node_op of+ "null" ->+ when (not $ elem nodeid blacklist) $+ mkNode nodeid (#label := _node_name .& #shape := GV.Ellipse .& #fillcolor := color0 .& Nil) "Convolution" -> do let attr = fromJust $ _node_attrs krnl = formatTuple (fromJust $ M.lookup "kernel" attr) strd = formatTuple (fromMaybe "1" $ M.lookup "stride" attr) nflt = fromJust $ M.lookup "num_filter" attr- lbl = printf "Convolution\n%s/%s, %s" krnl strd nflt- mkNode nodeid (#label := lbl .& #fillcolor := colors !! 1 .& Nil)+ lbl = sformat ("Convolution\n" % stext % "/" % stext % ", " % stext) krnl strd nflt+ mkNode nodeid (#label := lbl .& #fillcolor := color1 .& Nil) "FullyConnected" -> do let attr = fromJust $ _node_attrs hddn = fromJust $ M.lookup "num_hidden" attr- lbl = printf "FullyConnected\n%s" hddn- mkNode nodeid (#label := lbl .& #fillcolor := colors !! 1 .& Nil)+ lbl = sformat ("FullyConnected\n" % stext) hddn+ mkNode nodeid (#label := lbl .& #fillcolor := color1 .& Nil) "BatchNorm" ->- mkNode nodeid (#label := "batchNorm" .& #fillcolor := colors !! 3 .& Nil)+ mkNode nodeid (#label := "batchNorm" .& #fillcolor := color3 .& Nil) "Activation" -> do let attr = fromJust $ _node_attrs actt = fromJust $ M.lookup "act_type" attr- lbl = printf "Activation\n%s" actt- mkNode nodeid (#label := lbl .& #fillcolor := colors !! 2 .& Nil)+ lbl = sformat ("Activation\n" % stext) actt+ mkNode nodeid (#label := lbl .& #fillcolor := color2 .& Nil) "LeakyReLU" -> do let attr = fromJust $ _node_attrs actt = fromJust $ M.lookup "act_type" attr- lbl = printf "LeakyReLU\n%s" actt- mkNode nodeid (#label := lbl .& #fillcolor := colors !! 2 .& Nil)+ lbl = sformat ("LeakyReLU\n" % stext) actt+ mkNode nodeid (#label := lbl .& #fillcolor := color2 .& Nil) "Pooling" -> do let attr = fromJust $ _node_attrs poot = fromJust $ M.lookup "pool_type" attr krnl = formatTuple (fromJust $ M.lookup "kernel" attr) strd = formatTuple (fromMaybe "1" $ M.lookup "stride" attr)- lbl = printf "Pooling\n%s, %s/%s" poot krnl strd- mkNode nodeid (#label := lbl .& #fillcolor := colors !! 4 .& Nil)- "Concat" -> - mkNode nodeid (#label := "Concat" .& #fillcolor := colors !! 5 .& Nil)+ lbl = sformat ("Pooling\n" % stext % ", " % stext % "/" % stext) poot krnl strd+ mkNode nodeid (#label := lbl .& #fillcolor := color4 .& Nil)+ "Concat" ->+ mkNode nodeid (#label := "Concat" .& #fillcolor := color5 .& Nil) "Flatten" ->- mkNode nodeid (#label := "Flatten" .& #fillcolor := colors !! 5 .& Nil)+ mkNode nodeid (#label := "Flatten" .& #fillcolor := color5 .& Nil) "Reshape" ->- mkNode nodeid (#label := "Reshape" .& #fillcolor := colors !! 5 .& Nil)+ mkNode nodeid (#label := "Reshape" .& #fillcolor := color5 .& Nil) "Softmax" ->- mkNode nodeid (#label := "Softmax" .& #fillcolor := colors !! 6 .& Nil)+ mkNode nodeid (#label := "Softmax" .& #fillcolor := color6 .& Nil) "Custom" -> do let attr = fromJust $ _node_attrs lbl = fromJust $ M.lookup "op_type" attr- mkNode nodeid (#label := lbl .& #fillcolor := colors !! 7 .& Nil)+ mkNode nodeid (#label := lbl .& #fillcolor := color7 .& Nil) _ ->- mkNode nodeid (#label := _node_name .& #fillcolor := colors !! 7 .& Nil)+ mkNode nodeid (#label := _node_name .& #fillcolor := color7 .& Nil) mkEdge_ blacklist (tid, tnode) = do let op = _node_op tnode -- name = _node_name tnode- case op of + case op of "null" -> return () _ -> forM_ (_node_inputs tnode) $ \(sid:_) -> do- when (not $ elem sid blacklist) $ + when (not $ elem sid blacklist) $ GVM.edge tid sid [GV.Dir GV.Back, GV.ArrowTail GV.vee] - colors = catMaybes $ map color ["#8dd3c7", "#fb8072", "#ffffb3", - "#bebada", "#80b1d3", "#fdb462", - "#b3de69", "#fccde5"]+ [ color0, color1, color2, color3, color4, color5, color6, color7 ] =+ catMaybes $ map color ["#8dd3c7", "#fb8072", "#ffffb3",+ "#bebada", "#80b1d3", "#fdb462",+ "#b3de69", "#fccde5"] - _like sfx node = T.isSuffixOf sfx (T.pack $ _node_name node)+ _like sfx node = T.isSuffixOf sfx (_node_name node) -type instance ParameterList "graphviz_node" = - '[ '("label", 'AttrOpt String),+type instance ParameterList "graphviz_node" t =+ '[ '("label", 'AttrOpt Text), '("shape", 'AttrOpt GV.Shape), '("fixedsize", 'AttrOpt Bool),- '("fillcolor", 'AttrOpt GV.Color), - '("width", 'AttrOpt Double), - '("height", 'AttrOpt Double), + '("fillcolor", 'AttrOpt GV.Color),+ '("width", 'AttrOpt Double),+ '("height", 'AttrOpt Double), '("style", 'AttrOpt GV.Style) ] -mkNode :: (Fullfilled "graphviz_node" args)- => Int -> ArgsHMap "graphviz_node" args -> GVM.DotM Int ()+mkNode :: (Fullfilled "graphviz_node" () args)+ => Int -> ArgsHMap "graphviz_node" () args -> GVM.DotM Int () mkNode nodeid args = GVM.node nodeid attrs where shp = GV.Shape $ fromMaybe GV.BoxShape $ args !? #shape fxs = GV.FixedSize $ if fromMaybe True (args !? #fixedsize)- then GV.SetNodeSize + then GV.SetNodeSize else GV.GrowAsNeeded wdt = GV.Width $ fromMaybe 1.3 $ args !? #width hgt = GV.Height $ fromMaybe 0.8034 $ args !? #height sty = GV.style $ fromMaybe GV.filled $ args !? #style mfc = maybeToList $ GV.FillColor . GV.toColorList . (:[]) <$> (args !? #fillcolor)- lbl = maybeToList $ GV.textLabel . T.pack <$> (args !? #label)+ lbl = maybeToList $ GV.textLabel . TL.fromStrict <$> (args !? #label) attrs = [shp, fxs, wdt, hgt, sty] ++ lbl ++ mfc color :: String -> Maybe GV.Color@@ -174,11 +171,14 @@ return $ GV.RGB r g b color _ = Nothing -formatTuple :: String -> String-formatTuple str - | [((a,b),"")] <- (reads :: ReadS (Int,Int)) str = printf "%dx%d" a b- | [([a,b],"")] <- (reads :: ReadS [Int]) str = printf "%dx%d" a b+formatTuple :: Text -> Text+formatTuple str+ | Just (a,b) <- readMaybe sstr = sformat pf (a :: Int) (b :: Int)+ | Just [a,b] <- readMaybe sstr = sformat pf (a :: Int) b | otherwise = str+ where+ sstr = T.unpack str+ pf = int % "x" % int data Exc = CannotDecodeJSONofSymbol deriving (Show, Typeable)
+ src/MXNet/NN/Utils/Render.hs view
@@ -0,0 +1,66 @@+{-# LANGUAGE OverloadedLists #-}+module MXNet.NN.Utils.Render where++import Codec.Picture.Types+import Data.Array.Repa ((:.) (..), Array, DIM3, U, Z (..),+ extent, toUnboxed)+import qualified Graphics.Rasterific as G+import qualified Graphics.Rasterific.Texture as G (uniformTexture)+import Graphics.Text.TrueType (Font)+import RIO+import qualified RIO.Text as T+import qualified RIO.Vector.Storable as V++import MXNet.Base+++render :: (G.RenderablePixel px, ColorConvertible Pixel8 px) => Int -> Int -> G.Drawing px () -> Image px+render width height = G.renderDrawing width height (promotePixel (0 :: Pixel8))+++drawImage :: Image px -> G.Drawing px ()+drawImage img = G.drawImage img 0 (G.V2 0 0)+++drawBox :: px -> Float -> Float -> Float -> Float -> Float -> Maybe (Text, Font, Float, px) -> G.Drawing px ()+drawBox color stroke_width x0 y0 x1 y1 label = do+ G.withTexture (G.uniformTexture color) $+ G.stroke stroke_width G.JoinRound (G.CapRound, G.CapRound) $+ G.rectangle (G.V2 x0 y0) (x1 - x0) (y1 - y0)++ case label of+ Nothing -> return ()+ Just (text, font, size, text_color) ->+ G.withTexture (G.uniformTexture text_color) $+ let text_str = T.unpack text+ in G.printTextAt font (G.PointSize size) (G.V2 (x0+2) (y0+size+2)) text_str+++data NotConvertible = NotConvertible (NonEmpty Int)+ deriving Show+instance Exception NotConvertible+++imageFromNDArray :: (ColorConvertible PixelRGB8 px, HasCallStack)+ => NDArray Float -> IO (Image px)+imageFromNDArray array = do+ shape <- ndshape array+ case shape of+ [height, width, 3] -> do+ vec <- toVector array+ let img = Image width height (V.map floor vec) :: Image PixelRGB8+ return $ promoteImage img+ _ -> throwM $ NotConvertible shape+++imageFromRepa :: (ColorConvertible PixelRGB8 px, HasCallStack)+ => Array U DIM3 Float -> Image px+imageFromRepa array | c == 3 = promoted+ | otherwise = impureThrow (NotConvertible shape)+ where+ Z :. h :. w :. c = extent array+ shape = [h, w, c]+ vec = V.convert $ toUnboxed array+ image = Image w h (V.map floor vec) :: Image PixelRGB8+ promoted = promoteImage image+
+ src/MXNet/NN/Utils/Repa.hs view
@@ -0,0 +1,144 @@+{-# LANGUAGE ExplicitForAll #-}+{-# LANGUAGE Rank2Types #-}+module MXNet.NN.Utils.Repa where++import RIO+import RIO.List (splitAt)+import RIO.List.Partial (last)+import qualified RIO.Text as T (pack)+import qualified RIO.Vector.Boxed as V+import qualified RIO.Vector.Boxed.Partial as V (tail, foldl1')+import qualified RIO.Vector.Unboxed as VU+import qualified RIO.Vector.Unboxed.Partial as VU (maxIndex)+import Control.Exception (throw)+import Control.Lens+import Data.Array.Repa (Shape, Array, U, DIM1, DIM2, DIM3, DIM4, All(..), Z(..), (:.)(..), extent, toUnboxed)+import qualified Data.Array.Repa as Repa+import Text.PrettyPrint.Leijen.Text (Pretty(..), (<+>), textStrict)++newtype PrettyArray u s e = PrettyArray (Array u s e)+instance (Pretty e, VU.Unbox e, Shape d) => Pretty (PrettyArray U d e) where+ pretty (PrettyArray arr) = textStrict (T.pack $ Repa.showShape $ extent arr) <+> pretty (VU.toList $ toUnboxed arr)++class IxedReadOnly m where+ ixr :: Index m -> Fold m (IxValue m)++type instance Index (Array u sh a) = sh+type instance IxValue (Array u sh a) = a++instance (Repa.Source u a, Shape sh) => IxedReadOnly (Array u sh a) where+ ixr i f a+ | not (Repa.inShapeRange Repa.zeroDim (extent a) i) = pure a+ | otherwise = f (Repa.unsafeIndex a i) *> pure a++newtype ArrayFlatten u sh a = ArrayFlatten {getArray :: Array u sh a}++type instance Index (ArrayFlatten u sh a) = Int+type instance IxValue (ArrayFlatten u sh a) = a++(^#!) :: (Repa.Source u a, Shape sh, HasCallStack) => Array u sh a -> Int -> a+a ^#! i = ArrayFlatten a ^?! ixr i++instance (Repa.Source u a, Shape sh) => IxedReadOnly (ArrayFlatten u sh a) where+ ixr i f aflt@(ArrayFlatten a)+ | not (i >= 0 && i < Repa.size (extent a)) = pure aflt+ | otherwise = f (Repa.unsafeLinearIndex a i) *> pure aflt++expandDim :: (Shape sh, VU.Unbox e) => Int -> Array U sh e -> Array U (sh :. Int) e+expandDim axis arr | axis >=0 && axis < rank = Repa.computeS $ Repa.reshape shape_new arr+ | otherwise = error "Bad axis to expand."+ where+ shape = extent arr+ rank = Repa.rank shape+ (h, t) = splitAt (rank - axis) $ Repa.listOfShape shape+ shape_new = Repa.shapeOfList $ h ++ [1] ++ t+++vstack :: (Shape sh, VU.Unbox e) => V.Vector (Array U sh e) -> Array U sh e+-- alternative definition:+-- vstack = Repa.transpose . V.foldl1 (Repa.++) . V.map Repa.transpose+vstack arrs = Repa.fromUnboxed shape_new $ VU.concat $ V.toList $ V.map toUnboxed arrs+ where+ sumShape sh1 sh2 = let a1:r1 = reverse $ Repa.listOfShape sh1+ a2:r2 = reverse $ Repa.listOfShape sh2+ in if r1 == r2+ then Repa.shapeOfList $ reverse $ (a1+a2):r1+ else error "Cannot stack array because of incompatible shapes"+ shape_new = V.foldl1' sumShape $ V.map extent arrs+++vunstack :: (Unstackable sh, VU.Unbox e) => Array U sh e -> V.Vector (Array U (PredDIM sh) e)+vunstack arr = V.map (\i -> Repa.computeS $ Repa.slice arr (makeSliceAtAxis0 shape i)) range+ where+ shape = extent arr+ dim0 = last $ Repa.listOfShape shape+ range = V.enumFromN (0::Int) dim0++class (Shape sh,+ Shape (PredDIM sh),+ Repa.Slice (SliceAtAxis0 sh),+ Repa.FullShape (SliceAtAxis0 sh) ~ sh,+ Repa.SliceShape (SliceAtAxis0 sh) ~ PredDIM sh+ ) => Unstackable sh where+ type PredDIM sh+ type SliceAtAxis0 sh+ makeSliceAtAxis0 :: sh -> Int -> SliceAtAxis0 sh++instance Unstackable DIM2 where+ type PredDIM DIM2 = DIM1+ type SliceAtAxis0 DIM2 = Z:.Int:.All+ makeSliceAtAxis0 _ i = Z:.i:.All++instance Unstackable DIM3 where+ type PredDIM DIM3 = DIM2+ type SliceAtAxis0 DIM3 = Z:.Int:.All:.All+ makeSliceAtAxis0 (sh:._) i = makeSliceAtAxis0 sh i :. All++instance Unstackable DIM4 where+ type PredDIM DIM4 = DIM3+ type SliceAtAxis0 DIM4 = Z:.Int:.All:.All:.All+ makeSliceAtAxis0 (sh:._) i = makeSliceAtAxis0 sh i :. All++data ReshapeError = ReshapeMismatch (V.Vector Int) (V.Vector Int)+ | ReshapeTooManyMinusOne (V.Vector Int)+ deriving Show+instance Exception ReshapeError++reshapeEx :: (Shape sh1, Shape sh2, VU.Unbox e) => sh2 -> Array U sh1 e -> Array U sh2 e+reshapeEx shape arr = Repa.computeS $ Repa.reshape real_new_shape arr+ where+ old_shape = V.reverse $ V.fromList $ Repa.listOfShape $ extent arr+ new_shape = V.reverse $ V.fromList $ Repa.listOfShape shape+ shapeMismatch, tooManyN1 :: forall a. a+ shapeMismatch = throw (ReshapeMismatch new_shape old_shape)+ tooManyN1 = throw (ReshapeTooManyMinusOne new_shape)++ sizeEqual sh = V.product old_shape == V.product sh+ replaceZ i v | v == 0 = case old_shape V.!? i of+ Just v' -> v'+ Nothing -> shapeMismatch+ | otherwise = v+ new_shape_nz = V.imap replaceZ new_shape++ minus_n1s = V.elemIndices (-1) new_shape_nz+ filled_new_shape+ | V.null minus_n1s = if sizeEqual new_shape_nz then new_shape_nz else shapeMismatch+ | [s] <- V.toList minus_n1s = let (new_p1, new_p2) = V.splitAt s new_shape_nz+ in matchN1 new_p1 (V.tail new_p2) old_shape+ | otherwise = tooManyN1++ matchN1 sh1a sh1b sh2 | r == 0 = sh1a V.++ V.fromList [q] V.++ sh1b+ | otherwise = shapeMismatch+ where size1 = V.product $ sh1a V.++ sh1b+ size2 = V.product sh2+ (q, r) = size2 `divMod` size1++ real_new_shape = Repa.shapeOfList $ V.toList $ V.reverse filled_new_shape++argMax :: (VU.Unbox e, Ord e)+ => Array U DIM2 e -> V.Vector Int+--argMax overlaps =+-- let Z :. m :. n = extent overlaps+-- findMax row = VU.maxIndex $ toUnboxed $ Repa.computeS $ Repa.slice overlaps (Z :. row :. All)+-- in V.map findMax $ V.enumFromN (0 :: Int) m+argMax arr = V.map (VU.maxIndex . toUnboxed) (vunstack arr)