diff --git a/examples/mnist/DatasetVector.hs b/examples/mnist/DatasetVector.hs deleted file mode 100644 --- a/examples/mnist/DatasetVector.hs +++ /dev/null @@ -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 diff --git a/examples/mnist/Parse.hs b/examples/mnist/Parse.hs deleted file mode 100644 --- a/examples/mnist/Parse.hs +++ /dev/null @@ -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 diff --git a/examples/mnist/lenet.hs b/examples/mnist/lenet.hs deleted file mode 100644 --- a/examples/mnist/lenet.hs +++ /dev/null @@ -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) diff --git a/fei-nn.cabal b/fei-nn.cabal --- a/fei-nn.cabal +++ b/fei-nn.cabal @@ -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 + } diff --git a/src/MXNet/NN.hs b/src/MXNet/NN.hs --- a/src/MXNet/NN.hs +++ b/src/MXNet/NN.hs @@ -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 diff --git a/src/MXNet/NN/Callback.hs b/src/MXNet/NN/Callback.hs --- a/src/MXNet/NN/Callback.hs +++ b/src/MXNet/NN/Callback.hs @@ -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 + lr <- use (untag . mod_statistics . stat_last_lr) + lift . logInfo . display $ sformat ("") 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 diff --git a/src/MXNet/NN/DataIter/Class.hs b/src/MXNet/NN/DataIter/Class.hs --- a/src/MXNet/NN/DataIter/Class.hs +++ b/src/MXNet/NN/DataIter/Class.hs @@ -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 -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) diff --git a/src/MXNet/NN/DataIter/Conduit.hs b/src/MXNet/NN/DataIter/Conduit.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/DataIter/Conduit.hs @@ -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 diff --git a/src/MXNet/NN/DataIter/ConduitAsync.hs b/src/MXNet/NN/DataIter/ConduitAsync.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/DataIter/ConduitAsync.hs @@ -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) diff --git a/src/MXNet/NN/DataIter/Streaming.hs b/src/MXNet/NN/DataIter/Streaming.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/DataIter/Streaming.hs @@ -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 diff --git a/src/MXNet/NN/DataIter/Vec.hs b/src/MXNet/NN/DataIter/Vec.hs --- a/src/MXNet/NN/DataIter/Vec.hs +++ b/src/MXNet/NN/DataIter/Vec.hs @@ -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 diff --git a/src/MXNet/NN/EvalMetric.hs b/src/MXNet/NN/EvalMetric.hs --- a/src/MXNet/NN/EvalMetric.hs +++ b/src/MXNet/NN/EvalMetric.hs @@ -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 "" (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 "" (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 diff --git a/src/MXNet/NN/Initializer.hs b/src/MXNet/NN/Initializer.hs --- a/src/MXNet/NN/Initializer.hs +++ b/src/MXNet/NN/Initializer.hs @@ -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 diff --git a/src/MXNet/NN/Layer.hs b/src/MXNet/NN/Layer.hs --- a/src/MXNet/NN/Layer.hs +++ b/src/MXNet/NN/Layer.hs @@ -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 diff --git a/src/MXNet/NN/LrScheduler.hs b/src/MXNet/NN/LrScheduler.hs --- a/src/MXNet/NN/LrScheduler.hs +++ b/src/MXNet/NN/LrScheduler.hs @@ -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) diff --git a/src/MXNet/NN/Module.hs b/src/MXNet/NN/Module.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/Module.hs @@ -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) diff --git a/src/MXNet/NN/NDArray.hs b/src/MXNet/NN/NDArray.hs deleted file mode 100644 --- a/src/MXNet/NN/NDArray.hs +++ /dev/null @@ -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 diff --git a/src/MXNet/NN/Optimizer.hs b/src/MXNet/NN/Optimizer.hs --- a/src/MXNet/NN/Optimizer.hs +++ b/src/MXNet/NN/Optimizer.hs @@ -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 + diff --git a/src/MXNet/NN/Session.hs b/src/MXNet/NN/Session.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/Session.hs @@ -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 + diff --git a/src/MXNet/NN/TaggedState.hs b/src/MXNet/NN/TaggedState.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/TaggedState.hs @@ -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 "")) + diff --git a/src/MXNet/NN/Types.hs b/src/MXNet/NN/Types.hs --- a/src/MXNet/NN/Types.hs +++ b/src/MXNet/NN/Types.hs @@ -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 + diff --git a/src/MXNet/NN/Utils.hs b/src/MXNet/NN/Utils.hs --- a/src/MXNet/NN/Utils.hs +++ b/src/MXNet/NN/Utils.hs @@ -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 + diff --git a/src/MXNet/NN/Utils/GraphViz.hs b/src/MXNet/NN/Utils/GraphViz.hs --- a/src/MXNet/NN/Utils/GraphViz.hs +++ b/src/MXNet/NN/Utils/GraphViz.hs @@ -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) diff --git a/src/MXNet/NN/Utils/Render.hs b/src/MXNet/NN/Utils/Render.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/Utils/Render.hs @@ -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 + diff --git a/src/MXNet/NN/Utils/Repa.hs b/src/MXNet/NN/Utils/Repa.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/Utils/Repa.hs @@ -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)