fei-nn (empty) → 0.2.0
raw patch · 18 files changed
+1768/−0 lines, 18 filesdep +aesondep +attoparsecdep +attoparsec-binary
Dependencies added: aeson, attoparsec, attoparsec-binary, base, bytestring, containers, exceptions, fei-base, fei-nn, ghc-prim, graphviz, lens, mmorph, mtl, resourcet, template-haskell, text, time, transformers-base, unordered-containers, vector
Files
- LICENSE +29/−0
- examples/mnist/DatasetVector.hs +64/−0
- examples/mnist/Parse.hs +33/−0
- examples/mnist/lenet.hs +121/−0
- fei-nn.cabal +73/−0
- src/MXNet/NN.hs +422/−0
- src/MXNet/NN/Callback.hs +69/−0
- src/MXNet/NN/DataIter/Class.hs +41/−0
- src/MXNet/NN/DataIter/Vec.hs +42/−0
- src/MXNet/NN/EvalMetric.hs +136/−0
- src/MXNet/NN/Initializer.hs +61/−0
- src/MXNet/NN/Layer.hs +100/−0
- src/MXNet/NN/LrScheduler.hs +76/−0
- src/MXNet/NN/NDArray.hs +19/−0
- src/MXNet/NN/Optimizer.hs +133/−0
- src/MXNet/NN/Types.hs +99/−0
- src/MXNet/NN/Utils.hs +65/−0
- src/MXNet/NN/Utils/GraphViz.hs +185/−0
+ LICENSE view
@@ -0,0 +1,29 @@+BSD 3-Clause License++Copyright (c) 2018, Jiasen Wu+All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++* Redistributions of source code must retain the above copyright notice, this+ list of conditions and the following disclaimer.++* Redistributions in binary form must reproduce the above copyright notice,+ this list of conditions and the following disclaimer in the documentation+ and/or other materials provided with the distribution.++* Neither the name of the copyright holder nor the names of its+ contributors may be used to endorse or promote products derived from+ this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ examples/mnist/DatasetVector.hs view
@@ -0,0 +1,64 @@+module DatasetVector where + +import MXNet.Base (Symbol, NDArray, makeNDArray, contextCPU) +import Data.Typeable +import Control.Monad.Trans.Resource (MonadResource(..), MonadThrow(..)) +import Control.Monad.IO.Class (liftIO) +import Control.Monad (liftM2) +import Control.Exception.Base +import qualified Data.Vector as V +import qualified Data.Vector.Storable as VS +import qualified Data.ByteString as BS +import Data.Attoparsec.ByteString as AP + +import MXNet.NN.DataIter.Class +import MXNet.NN.DataIter.Vec +import Parse + +type SymbolF = Symbol Float +type ArrayF = NDArray Float + +loadTrainingData :: (MonadResource m, MonadThrow m) => m (DatasetVector (ArrayF, ArrayF)) +loadTrainingData = do + v1 <- batch 128 <$> sourceImages "examples/data/train-images-idx3-ubyte" + v2 <- batch 128 <$> sourceLabels "examples/data/train-labels-idx1-ubyte" + liftIO $ liftM2 zipD (mapMD cImageToNDArray v1) (mapMD cLabelToNDArray v2) + +loadTestingData :: (MonadResource m, MonadThrow m) => m (DatasetVector (ArrayF, ArrayF)) +loadTestingData = do + v1 <- batch 1 <$> sourceImages "examples/data/t10k-images-idx3-ubyte" + v2 <- batch 1 <$> sourceLabels "examples/data/t10k-labels-idx1-ubyte" + liftIO $ liftM2 zipD (mapMD cImageToNDArray v1) (mapMD cLabelToNDArray v2) + +sourceImages :: (MonadResource m, MonadThrow m) => FilePath -> m (DatasetVector Image) +sourceImages = parseFile $ do + HeaderImg n w h <- header + count n (image w h) + +sourceLabels :: (MonadResource m, MonadThrow m) => FilePath -> m (DatasetVector Label) +sourceLabels = parseFile $ do + HeaderLbl n <- header + count n label + +parseFile :: (MonadResource m, MonadThrow m) => Parser [a] -> FilePath -> m (DatasetVector a) +parseFile parser fp = do + content <- liftIO $ BS.readFile fp + case AP.parseOnly parser content of + Left msg -> throwM $ ParseError msg + Right rt -> return $ fromListD rt + +batch :: Int -> DatasetVector a -> DatasetVector (V.Vector a) +batch n (DatasetVector vec) = (DatasetVector $ walk n vec) + where + walk n vec = V.unfoldr (\v -> if V.null v then Nothing else Just (V.splitAt n v)) vec + +mapMD f (DatasetVector vec) = DatasetVector <$> V.mapM f vec + +cImageToNDArray :: V.Vector Image -> IO ArrayF +cImageToNDArray dat = makeNDArray [V.length dat, 1, 28, 28] contextCPU (VS.concat $ V.toList dat) +cLabelToNDArray :: V.Vector Label -> IO ArrayF +cLabelToNDArray dat = makeNDArray [V.length dat] contextCPU (V.convert $ V.map fromIntegral dat) + +data Exc = ParseError String + deriving (Show, Typeable) +instance Exception Exc
+ examples/mnist/Parse.hs view
@@ -0,0 +1,33 @@+module Parse where++import Data.Attoparsec.ByteString as AP+import Data.Attoparsec.Binary as AP+import qualified Data.ByteString.Internal as BS+import qualified Data.Vector.Storable as SV++type Image = SV.Vector Float+type Label = Int++data Header = HeaderImg Int Int Int+ | HeaderLbl Int++header :: AP.Parser Header+header = do+ mc <- AP.anyWord32be+ case mc of+ 0x00000803 -> do + [d1,d2,d3] <- AP.count 3 AP.anyWord32be+ return $ HeaderImg (fromIntegral d1) (fromIntegral d2) (fromIntegral d3)+ 0x00000801 -> do + d1 <- AP.anyWord32be+ return $ HeaderLbl (fromIntegral d1)+ _ -> fail "Header type not recognised"++image :: Int -> Int -> AP.Parser Image+image w h = do+ BS.PS fp ofs len <- AP.take (w*h)+ let vw = SV.unsafeFromForeignPtr fp ofs len+ return $ SV.map ((/255) . fromIntegral) vw++label :: AP.Parser Label+label = fromIntegral <$> AP.anyWord8
+ examples/mnist/lenet.hs view
@@ -0,0 +1,121 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedLabels #-}+module Main where++import Control.Monad (forM_, void)+import qualified Data.Vector.Storable as SV+import Control.Monad.IO.Class+import Control.Monad.Trans.Resource+import System.IO (hFlush, stdout)+import qualified Data.HashMap.Strict as M++import MXNet.Base hiding (zeros)+import qualified MXNet.Base.Operators.NDArray as A+import MXNet.NN+import MXNet.NN.DataIter.Class+import qualified MXNet.NN.Utils.GraphViz as GV++import DatasetVector++-- # first conv+-- conv1 = mx.symbol.Convolution(data=data, kernel=(5,5), num_filter=20)+-- tanh1 = mx.symbol.Activation(data=conv1, act_type="tanh")+-- pool1 = mx.symbol.Pooling(data=tanh1, pool_type="max", kernel=(2,2), stride=(2,2))+-- # second conv+-- conv2 = mx.symbol.Convolution(data=pool1, kernel=(5,5), num_filter=50)+-- tanh2 = mx.symbol.Activation(data=conv2, act_type="tanh")+-- pool2 = mx.symbol.Pooling(data=tanh2, pool_type="max", kernel=(2,2), stride=(2,2))+-- # first fullc+-- flatten = mx.symbol.Flatten(data=pool2)+-- fc1 = mx.symbol.FullyConnected(data=flatten, num_hidden=500)+-- tanh3 = mx.symbol.Activation(data=fc1, act_type="tanh")+-- # second fullc+-- fc2 = mx.symbol.FullyConnected(data=tanh3, num_hidden=num_classes)+-- # loss+-- lenet = mx.symbol.SoftmaxOutput(data=fc2, name='softmax')++neural :: IO SymbolF+neural = do+ x <- variable "x"+ y <- variable "y"++ v1 <- convolution "conv1" (#data := x .& #kernel := [5,5] .& #num_filter := 20 .& Nil)+ a1 <- activation "conv1-a" (#data := v1 .& #act_type := #tanh .& Nil)+ p1 <- pooling "conv1-p" (#data := a1 .& #kernel := [2,2] .& #pool_type := #max .& Nil)++ v2 <- convolution "conv2" (#data := p1 .& #kernel := [5,5] .& #num_filter := 50 .& Nil)+ a2 <- activation "conv2-a" (#data := v2 .& #act_type := #tanh .& Nil)+ p2 <- pooling "conv2-p" (#data := a2 .& #kernel := [2,2] .& #pool_type := #max .& Nil)++ fl <- flatten "flatten" (#data := p2 .& Nil)++ v3 <- fullyConnected "fc1" (#data := fl .& #num_hidden := 500 .& Nil)+ a3 <- activation "fc1-a" (#data := v3 .& #act_type := #tanh .& Nil)++ v4 <- fullyConnected "fc2" (#data := a3 .& #num_hidden := 10 .& Nil)+ a4 <- softmaxoutput "softmax" (#data := v4 .& #label := y .& Nil)+ return $ Symbol a4++range :: Int -> [Int]+range = enumFromTo 1++default_initializer :: Initializer Float+default_initializer name shp@[_] = zeros name shp+default_initializer name shp@[_,_] = xavier 2.0 XavierGaussian XavierIn name shp+default_initializer name shp = normal 0.1 name shp++main :: IO ()+main = do+ -- call mxListAllOpNames can ensure the MXNet itself is properly initialized+ -- i.e. MXNet operators are registered in the NNVM+ _ <- mxListAllOpNames+ net <- neural+ -- GV.dotPlot net GV.Png "lenet"+ sess <- initialize net $ Config { + _cfg_data = M.singleton "x" [1,28,28],+ _cfg_label = ["y"],+ _cfg_initializers = M.empty,+ _cfg_default_initializer = default_initializer,+ _cfg_context = contextCPU+ }+ optimizer <- makeOptimizer SGD'Mom (Const 0.0002) (#momentum := 0.9 .& #wd := 0.0001 .& Nil)++ runResourceT $ train sess $ do ++ trainingData <- loadTrainingData+ testingData <- loadTestingData++ liftIO $ putStrLn $ "[Train] "+ forM_ (range 5) $ \ind -> do+ liftIO $ putStrLn $ "iteration " ++ show ind+ metric <- newMetric "train" (CrossEntropy "y")+ void $ forEachD_ni trainingData $ \((t,i), (x, y)) -> do+ eval <- format metric+ liftIO $ putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show t ++ " " ++ eval+ liftIO $ hFlush stdout+ fitAndEval optimizer (M.fromList [("x", x), ("y", y)]) metric+ liftIO $ putStrLn ""+ + liftIO $ putStrLn $ "[Test] "+ result <- forEachD_ni testingData $ \((t,i), (x, y)) -> do + liftIO $ do + putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show t+ hFlush stdout+ [y'] <- forwardOnly (M.fromList [("x", Just x), ("y", Nothing)])+ ind1 <- liftIO $ toVector y+ ind2 <- liftIO $ argmax y' >>= toVector+ return (ind1, ind2)+ liftIO $ putStr "\r\ESC[K"++ let (ls,ps) = unzip result+ ls_unbatched = mconcat ls+ ps_unbatched = mconcat ps+ total_test_items = SV.length ls_unbatched+ correct = SV.length $ SV.filter id $ SV.zipWith (==) ls_unbatched ps_unbatched+ liftIO $ putStrLn $ "Accuracy: " ++ show correct ++ "/" ++ show total_test_items+ + where+ argmax :: ArrayF -> IO ArrayF+ argmax (NDArray ys) = NDArray . head <$> A.argmax (#data := ys .& #axis := Just 1 .& Nil)
+ fei-nn.cabal view
@@ -0,0 +1,73 @@+name: fei-nn+version: 0.2.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-file: LICENSE+author: Jiasen Wu+maintainer: jiasenwu@hotmail.com+copyright: Copyright: (c) 2018 Jiasen Wu+category: Machine Learning, AI+build-type: Simple+cabal-version: 1.24++Library+ exposed-modules: MXNet.NN+ MXNet.NN.NDArray+ MXNet.NN.Types+ MXNet.NN.Utils+ MXNet.NN.Utils.GraphViz+ MXNet.NN.Layer+ MXNet.NN.Optimizer+ MXNet.NN.LrScheduler+ MXNet.NN.EvalMetric+ MXNet.NN.Initializer+ MXNet.NN.Callback+ MXNet.NN.DataIter.Class+ MXNet.NN.DataIter.Vec+ other-modules:+ hs-source-dirs: src+ ghc-options: -Wall+ default-language: Haskell2010+ default-extensions: GADTs,+ TypeFamilies,+ OverloadedLabels+ 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+ , fei-base+ , fei-nn
+ src/MXNet/NN.hs view
@@ -0,0 +1,422 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE FlexibleContexts #-}+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.Optimizer,+ module MXNet.NN.LrScheduler,+ module MXNet.NN.EvalMetric,+ module MXNet.NN.Initializer,+ module MXNet.NN.Layer,+ module MXNet.NN.Callback,+) 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)++-- | 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++ arg_tensors <- M.traverseWithKey (initI placeholders spec2 dinit) arg_with_shp+ aux_tensors <- M.traverseWithKey (initA dinit) aux_with_shp++ 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+ }+ 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++-- | 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++ 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++ 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)++ 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++-- | 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++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+ )++wait :: (MonadIO m, DType a) => NDArray a -> TrainM a m ()+wait (NDArray hdl) = liftIO $ mxNDArrayWaitToRead hdl++getContext :: Monad m => TrainM a m Context+getContext = use sess_context++-- | 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
+ src/MXNet/NN/Callback.hs view
@@ -0,0 +1,69 @@+module MXNet.NN.Callback where++import Control.Monad.State.Strict (lift)+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++-- | Learning rate+data DumpLearningRate = DumpLearningRate++instance CallbackClass DumpLearningRate where+ endOfBatch _ _ _ = do+ lr <- lift $ use stat_last_lr+ liftIO $ do+ putStr $ printf "<lr: %0.6f> " lr++-- | Throughput +data DumpThroughputEpoch = DumpThroughputEpoch {+ _tp_begin_time :: IORef UTCTime,+ _tp_end_time :: IORef UTCTime,+ _tp_total_sample :: IORef Int+}++instance CallbackClass DumpThroughputEpoch where+ begOfBatch _ n (DumpThroughputEpoch _ _ totalRef) = do+ liftIO $ modifyIORef totalRef (+n)+ begOfEpoch _ _ (DumpThroughputEpoch tt1Ref _ _) =+ liftIO $ getCurrentTime >>= writeIORef tt1Ref+ endOfEpoch _ _ (DumpThroughputEpoch _ tt2Ref _) = do+ liftIO $ getCurrentTime >>= writeIORef tt2Ref+ endOfVal _ _ (DumpThroughputEpoch tt1Ref tt2Ref totalRef) = liftIO $ 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++dumpThroughputEpoch :: IO Callback+dumpThroughputEpoch = do+ t0 <- getCurrentTime+ r0 <- newIORef t0+ r1 <- newIORef t0+ r2 <- newIORef 0+ return $ Callback $ DumpThroughputEpoch r0 r1 r2++-- | Checkpoint+data Checkpoint = Checkpoint String++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
+ src/MXNet/NN/DataIter/Class.hs view
@@ -0,0 +1,41 @@+{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleInstances #-} +module MXNet.NN.DataIter.Class where + +import GHC.Exts (Constraint) + +-- | Constraints on Dataset and the monad where the operation shall be ran. +type family DatasetConstraint (d :: * -> *) (m :: * -> *) :: Constraint + +-- | Abstract Dataset type class. +-- Available instances include 'LVec' and mxnet data-iters in package <https://github.com/pierric/mxnet-dataiter mxnet-dataiter> +class Dataset (d :: * -> *) where + -- | Create Dataset from `[]`. + -- note that depending on the instance, it may or may not work with infinitive list. + fromListD :: [e] -> d e + -- | Zip two Datasets + zipD :: d e1 -> d e2 -> d (e1, e2) + -- | Get number of elements + sizeD :: (DatasetConstraint d m, Monad m) => d e -> m Int + -- | Apply a function on each element of Dataset + forEachD :: (DatasetConstraint d m, Monad m) => d e -> (e -> m a) -> m [a] + + -- | Apply a function on each element of Dataset together with the element's index. + -- Note that the default implmentation assumes the Dataset can be created from a infinitive list. + forEachD_i :: (DatasetConstraint d m, Monad m) => d e -> ((Int, e) -> m a) -> m [a] + forEachD_i dat = forEachD (zipD (fromListD [1..]) dat) + + -- | Apply a function on each element of Dataset together with the total number of elements and the element's index. + forEachD_ni :: (DatasetConstraint d m, Monad m) => d e -> (((Int, Int), e) -> m a) -> m [a] + forEachD_ni dat proc = do + n <- sizeD dat + forEachD ((fromListD (replicate n n) `zipD` fromListD [1..n]) `zipD` dat) proc + + foldD :: (DatasetConstraint d m, Monad m) => (a -> e -> m a) -> a -> d e -> m a + + takeD :: Int -> d e -> d e + + +class DatasetProp (d :: * -> *) e where + -- | Get the batch size of the dataset + batchSizeD :: (DatasetConstraint d m, Monad m) => d e -> m (Maybe Int)
+ src/MXNet/NN/DataIter/Vec.hs view
@@ -0,0 +1,42 @@+{-# LANGUAGE MultiParamTypeClasses #-}+{-# 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 MXNet.NN.DataIter.Class+import MXNet.NN.Types+import MXNet.Base (NDArray, DType, ndshape)++newtype DatasetVector a = DatasetVector { _dsv_unwrap :: Vector a }+++type instance DatasetConstraint DatasetVector m = MonadIO m++instance Dataset DatasetVector where+ fromListD = DatasetVector . V.fromList+ zipD v1 v2 = DatasetVector $ V.zip (_dsv_unwrap v1) (_dsv_unwrap v2)+ sizeD = return . V.length . _dsv_unwrap+ forEachD dat func = V.toList <$> V.forM (_dsv_unwrap dat) func+ 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++instance DType a => DatasetProp DatasetVector (NDArray a) where+ batchSizeD (DatasetVector dat) = liftIO $ do+ batch_size : _ <- 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+ return $ if batch_size1 /= batch_size2+ then Nothing+ else Just batch_size1
+ src/MXNet/NN/EvalMetric.hs view
@@ -0,0 +1,136 @@+{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE DataKinds #-} +{-# 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 MXNet.Base +import qualified MXNet.Base.Operators.NDArray as A +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 + + +-- | Basic evaluation - accuracy +data Accuracy a = Accuracy String + +instance EvalMetricMethod Accuracy where + data MetricData Accuracy a = AccuracyData String String (IORef Int) (IORef Int) + newMetric phase (Accuracy label) = 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 + + 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 + s <- liftIO $ readIORef sumRef + n <- liftIO $ readIORef cntRef + return $ printf "<Accuracy: %0.2f>" (100 * fromIntegral s / fromIntegral n :: Float) + +-- | Basic evaluation - cross entropy +data CrossEntropy a = CrossEntropy String + +copyTo :: DType a => NDArray a -> NDArray a -> IO () +copyTo (NDArray dst) (NDArray src) = A._copyto_upd [dst] (#data := src .& Nil) + +instance EvalMetricMethod CrossEntropy where + data MetricData CrossEntropy a = CrossEntropyData String String (IORef Int) (IORef Float) + newMetric phase (CrossEntropy label) = do + a <- liftIO $ newIORef 0 + b <- liftIO $ newIORef 0 + return $ CrossEntropyData phase label a b + -- | evaluate the log-loss. + -- preds is of shape (batch_size, num_category), each element along the second dimension gives the probability of the category. + -- label is of shape (batch_size,), each element gives the category number. + evaluate (CrossEntropyData phase label cntRef sumRef) bindings [output] = do + liftIO $ compute output (bindings M.! label) + s <- liftIO $ readIORef sumRef + n <- liftIO $ readIORef cntRef + let loss = realToFrac s / fromIntegral n + return $ M.singleton (phase ++ "_loss") loss + where + compute preds lbl@(NDArray labelHandle) = do + shp1 <- ndshape preds + shp2 <- ndshape lbl + when (length shp1 /= 2 || length shp2 /= 1 || head shp1 /= head shp2) (throwM $ MismatchedShapeInEval shp1 shp2) + -- before call pick, we have to make sure preds and label + -- are in the same context + NDArray preds_may_copy <- do + c1 <- context preds + c2 <- context lbl + if c1 == c2 + then return preds + else do + preds_shap <- ndshape preds + preds_copy <- makeEmptyNDArray preds_shap c2 + copyTo preds_copy preds + return preds_copy + [predprj] <- A.pick (#data := preds_may_copy .& #index := labelHandle .& Nil) + [predlog] <- A.log (#data := predprj .& Nil) + loss <- A.sum (#data := predlog .& Nil) >>= toVector . NDArray . head + modifyIORef sumRef (+ (negate $ loss SV.! 0)) + modifyIORef cntRef (+ head shp1) + format (CrossEntropyData _ _ cntRef sumRef) = liftIO $ do + s <- liftIO $ readIORef sumRef + n <- liftIO $ readIORef cntRef + return $ printf "<CrossEntropy: %0.3f>" (realToFrac s / fromIntegral n :: Float) + +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 "" + +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 + return $ M.union m1 m2 + format (MCompositeData a as) = do + s1 <- format a + s2 <- format as + return $ s1 ++ " " ++ s2 + +infixr 9 :*
+ src/MXNet/NN/Initializer.hs view
@@ -0,0 +1,61 @@+{-# 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 MXNet.NN.Types+import MXNet.NN.Utils++empty :: DType a => Initializer a+empty _ shp cxt = makeEmptyNDArray shp cxt++zeros :: DType a => Initializer a+zeros = constant 0++ones :: DType a => Initializer a+ones = constant 1++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"]) + => Float -> Initializer a+uniform sca _ shp cxt = NDArray . head <$> (A._random_uniform + ( #low := (-sca) + .& #high := sca+ .& #shape := shp+ .& #ctx := formatContext cxt+ .& #dtype := EnumType (typename (undefined :: a))+ .& Nil))++normal :: forall a. (DType a, HasEnum (DTypeName a) '["None", "float16" ,"float32", "float64"]) + => Float -> Initializer a+normal sigma _ shp cxt = NDArray . head <$> (A._random_normal+ ( #loc := (0 :: Float)+ .& #scale := sigma+ .& #shape := shp+ .& #ctx := formatContext cxt+ .& #dtype := EnumType (typename (undefined :: a))+ .& Nil))++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"
+ src/MXNet/NN/Layer.hs view
@@ -0,0 +1,100 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE UndecidableInstances #-}++module MXNet.NN.Layer (+ variable,+ convolution,+ fullyConnected,+ pooling,+ activation,+ softmaxoutput,+ batchnorm,+ cast,+ plus,+ flatten,+ identity,+ dropout,+ reshape,+) where++import MXNet.Base+import qualified MXNet.Base.Operators.Symbol as S++variable :: String -> IO SymbolHandle+variable = mxSymbolCreateVariable++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 + then+ S._Convolution name (#weight := w .& args)+ else+ S._Convolution name (#bias := b .& #weight := w .& args)++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 + then+ S._FullyConnected name (#weight := w .& args)+ else+ S._FullyConnected name (#bias := b .& #weight := w .& args)++-- 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++activation :: HasArgs "_Activation(symbol)" args '["data", "act_type"]+ => String -> ArgsHMap "_Activation(symbol)" args -> IO SymbolHandle+activation = S._Activation++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(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)++cast :: HasArgs "_Cast(symbol)" args '["data", "dtype"]+ => String -> ArgsHMap "_Cast(symbol)" args -> IO SymbolHandle+cast name args = S._Cast name args++plus :: HasArgs "elemwise_add(symbol)" args '["lhs", "rhs"]+ => String -> ArgsHMap "elemwise_add(symbol)" args -> IO SymbolHandle+plus = S.elemwise_add++flatten :: HasArgs "_Flatten(symbol)" args '["data"]+ => String -> ArgsHMap "_Flatten(symbol)" args -> IO SymbolHandle+flatten = S._Flatten++identity :: HasArgs "_copy(symbol)" args '["data"]+ => String -> ArgsHMap "_copy(symbol)" args -> IO SymbolHandle+identity = S._copy++-- 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++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+
+ src/MXNet/NN/LrScheduler.hs view
@@ -0,0 +1,76 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeOperators #-}+module MXNet.NN.LrScheduler where++import MXNet.Base.Spec.Operator+import Data.Maybe (fromMaybe)++class LrScheduler sch where+ getLR :: sch -> Int -> Float++instance LrScheduler Float where+ getLR = const++data Const = Const Float+instance LrScheduler Const where+ getLR (Const lr) = const lr++lrOfConst :: Float -> Const+lrOfConst = Const++data FactorScheduler = Factor Float Float Int Float+instance LrScheduler FactorScheduler where+ 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), + '("base", 'AttrOpt Float), '("stop", 'AttrOpt Float)]+ +lrOfFactor :: Fullfilled "lrOfFactor" args + => ArgsHMap "lrOfFactor" args -> FactorScheduler+lrOfFactor args = Factor base factor step stop+ where + factor = args ! #factor+ step = args ! #step+ base = fromMaybe 0.01 (args !? #base)+ stop = fromMaybe 1e-8 (args !? #stop)++data MultifactorScheduler = Multifactor Float Float [Int]+instance LrScheduler MultifactorScheduler where+ getLR (Multifactor base factor steps) nup = base * factor ^ (index nup steps)+ where+ index a bs = go a bs (0 :: Int)+ go _ [] n = n+ go a (b:bs) n = if b > a then n else go a bs (n+1)++type instance ParameterList "lrOfMultifactor" =+ '[ '("factor", 'AttrReq Float), '("steps", 'AttrReq [Int]), '("base", 'AttrOpt Float)]++lrOfMultifactor :: Fullfilled "lrOfMultifactor" args+ => ArgsHMap "lrOfMultifactor" args -> MultifactorScheduler+lrOfMultifactor args = Multifactor base factor steps+ where + factor = args ! #factor+ steps = args ! #steps+ base = fromMaybe 0.01 (args !? #base)++data PolyScheduler = Poly Float Float Int+instance LrScheduler PolyScheduler where+ getLR (Poly base power maxnup) nup =+ if nup < maxnup + then base * (1 - fromIntegral nup / fromIntegral maxnup) ** power+ else 0++type instance ParameterList "lrOfPoly" =+ '[ '("maxnup", 'AttrReq Int), '("power", 'AttrReq Float), '("base", 'AttrOpt Float)]++lrOfPoly :: Fullfilled "lrOfPoly" args+ => ArgsHMap "lrOfPoly" args -> PolyScheduler+lrOfPoly args = Poly base power maxnup+ where + maxnup = args ! #maxnup+ base = fromMaybe 0.01 (args !? #base)+ power = fromMaybe 2 (args !? #power)
+ src/MXNet/NN/NDArray.hs view
@@ -0,0 +1,19 @@+module MXNet.NN.NDArray where++import MXNet.Base+import qualified MXNet.Base.Operators.NDArray as I++reshape :: DType a => NDArray a -> [Int] -> IO (NDArray a)+reshape arr shp = do+ [hdl] <- I._Reshape (#data := unNDArray arr .& #shape := shp .& Nil)+ return $ NDArray hdl++transpose :: DType a => NDArray a -> [Int] -> IO (NDArray a)+transpose arr axes = do+ [hdl] <- I.transpose (#data := unNDArray arr .& #axes := axes .& Nil)+ return $ NDArray hdl++copy :: DType a => NDArray a -> NDArray a -> IO (NDArray a)+copy src dst = do+ I._copyto_upd [unNDArray dst] (#data := unNDArray src .& Nil)+ return dst
+ src/MXNet/NN/Optimizer.hs view
@@ -0,0 +1,133 @@+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE UndecidableInstances #-}++module MXNet.NN.Optimizer (+ Optimizer(..),+ OptimizerTag(..)+) where++import MXNet.Base hiding (Symbol)+import qualified MXNet.Base.Operators.NDArray as A++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)++-- | Abstract Optimizer type class+class Optimizer (opt :: * -> *) where+ data OptimizerTag opt :: *+ -- | Specific required arguments+ -- data ReqArgs opt :: *+ -- | 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)+ -- | run the optimizer with the input & expected tensor+ optimize :: (DType dtype, MonadIO m, MonadState Statistics m) + => opt dtype -- optimizer+ -> String -- symbol name to optimize+ -> NDArray dytpe -- parameter+ -> NDArray dtype -- gradient+ -> m ()++type family OptimizerSym (opt :: * -> *) :: Symbol+type family OptimizerCst (opt :: * -> *) dt (args :: [*]) :: Constraint++-- | 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++type instance OptimizerSym SGD_Opt = "sgd_update(ndarray)"+-- 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"]++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+ let lr = getLR sch nup+ stat_last_lr .= lr+ liftIO $ A.sgd_update_upd [weight] (+ #weight := weight .& + #grad := gradient .& + #lr := lr .& args)++-- | 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++type instance OptimizerSym SGD_Mom_Opt = "sgd_mom_update(ndarray)"+-- 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"]++instance Optimizer SGD_Mom_Opt where+ data OptimizerTag SGD_Mom_Opt = SGD'Mom+ makeOptimizer SGD'Mom sch args = do+ 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+ let lr = getLR sch nup+ 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)+ return mom+ Just (NDArray a) -> return a+ A.sgd_mom_update_upd [weight] (+ #weight := weight .& + #grad := gradient .& + #mom := momentum .& + #lr := lr .& args)++-- | 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++type instance OptimizerSym ADAM_Opt = "adam_update(ndarray)"+-- 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"]++instance Optimizer ADAM_Opt where+ data OptimizerTag ADAM_Opt = ADAM+ makeOptimizer ADAM sch args = do+ 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+ let lr = getLR sch nup+ 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)+ 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)
+ src/MXNet/NN/Types.hs view
@@ -0,0 +1,99 @@+{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE ExplicitForAll #-} +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 + +-- | 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) + +-- | 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 +-- method to initialize the variables. +-- +-- 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 +-- _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 +} + +-- | 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) + +-- | Possible exception in 'TrainM' +data Exc = MismatchedShapeOfSym String [Int] [Int] + | MismatchedShapeInEval [Int] [Int] + | NotAParameter String + | InvalidArgument String + | InferredShapeInComplete + | DatasetOfUnknownBatchSize + | LoadSessionInvalidTensorName String + | LoadSessionMismatchedTensorKind String + deriving (Show, Typeable) +instance Exception Exc + +makeLenses ''Config +makeLenses ''Statistics +makeLenses ''Session
+ src/MXNet/NN/Utils.hs view
@@ -0,0 +1,65 @@+{-# 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 MXNet.Base (+ Context(..), DType, NDArray(..), Symbol(..), + HMap(..), (.&), ArgOf(..),+ mxSymbolSaveToFile, mxNDArraySave, mxNDArrayLoad)+import MXNet.NN.Types+import qualified MXNet.Base.Operators.NDArray as A++-- | format a shape+formatShape :: [Int] -> String+formatShape shape = concat $ ["("] ++ intersperse "," (map show shape) ++ [")"]++-- | format a context+formatContext :: Context -> String+formatContext Context{..} = getDeviceName _device_type ++ "(" ++ show _device_id ++ ")"+ where + 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)++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+ liftIO $ do+ mxSymbolSaveToFile (filename ++ ".json") (unSymbol net)+ mxNDArraySave (filename ++ ".params") modelParams+ where+ getModelParam (key, ParameterI a _) = ("arg:" ++ key, unNDArray a)+ getModelParam (key, ParameterA a) = ("aux:" ++ key, unNDArray a)++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)
+ src/MXNet/NN/Utils/GraphViz.hs view
@@ -0,0 +1,185 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications, TypeOperators, DataKinds #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE RecordWildCards #-}+module MXNet.NN.Utils.GraphViz (+ dotPlot,+ 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 qualified Data.GraphViz.Attributes.Complete as GV+import qualified Data.GraphViz.Types.Monadic as GVM+import qualified Data.GraphViz.Types.Generalised as GVM++import MXNet.Base++-- The program `dot` must be found in the PATH.++dotPlot :: DType a => Symbol a -> 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)++instance FromJSON JSNode where+ parseJSON (Object v) = JSNode <$> v .: "op"+ <*> v .: "name"+ <*> v .:? "attrs"+ <*> v .: "inputs"+ parseJSON invalid = typeMismatch "JSNode" invalid++data JSGraph = JSGraph {+ _symbol_nodes :: [JSNode]+} deriving (Show)++instance FromJSON JSGraph where+ parseJSON (Object v) = JSGraph <$> v .: "nodes"+ parseJSON invalid = typeMismatch "JSGraph" invalid++-- 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+ js <- mxSymbolSaveToJSON sym+ auxnodes <- mxSymbolListAuxiliaryStates sym+ case eitherDecode $ pack 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) + 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)+ "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)+ "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)+ "BatchNorm" ->+ mkNode nodeid (#label := "batchNorm" .& #fillcolor := colors !! 3 .& 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)+ "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)+ "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)+ "Flatten" ->+ mkNode nodeid (#label := "Flatten" .& #fillcolor := colors !! 5 .& Nil)+ "Reshape" ->+ mkNode nodeid (#label := "Reshape" .& #fillcolor := colors !! 5 .& Nil)+ "Softmax" ->+ mkNode nodeid (#label := "Softmax" .& #fillcolor := colors !! 6 .& 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 := _node_name .& #fillcolor := colors !! 7 .& Nil)++ mkEdge_ blacklist (tid, tnode) = do+ let op = _node_op tnode+ -- name = _node_name tnode+ case op of + "null" -> return ()+ _ -> forM_ (_node_inputs tnode) $ \(sid:_) -> do+ 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"]++ _like sfx node = T.isSuffixOf sfx (T.pack $ _node_name node)++type instance ParameterList "graphviz_node" = + '[ '("label", 'AttrOpt String),+ '("shape", 'AttrOpt GV.Shape),+ '("fixedsize", 'AttrOpt Bool),+ '("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 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 + 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)+ attrs = [shp, fxs, wdt, hgt, sty] ++ lbl ++ mfc++color :: String -> Maybe GV.Color+color ['#',r1,r2,g1,g2,b1,b2] = do+ let dec = listToMaybe . map fst . readHex+ r <- dec [r1,r2]+ g <- dec [g1,g2]+ b <- dec [b1,b2]+ 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+ | otherwise = str++data Exc = CannotDecodeJSONofSymbol+ deriving (Show, Typeable)+instance Exception Exc