diff --git a/LICENSE b/LICENSE
new file mode 100644
--- /dev/null
+++ b/LICENSE
@@ -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.
diff --git a/examples/mnist/DatasetVector.hs b/examples/mnist/DatasetVector.hs
new file mode 100644
--- /dev/null
+++ b/examples/mnist/DatasetVector.hs
@@ -0,0 +1,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
diff --git a/examples/mnist/Parse.hs b/examples/mnist/Parse.hs
new file mode 100644
--- /dev/null
+++ b/examples/mnist/Parse.hs
@@ -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
diff --git a/examples/mnist/lenet.hs b/examples/mnist/lenet.hs
new file mode 100644
--- /dev/null
+++ b/examples/mnist/lenet.hs
@@ -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)
diff --git a/fei-nn.cabal b/fei-nn.cabal
new file mode 100644
--- /dev/null
+++ b/fei-nn.cabal
@@ -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
diff --git a/src/MXNet/NN.hs b/src/MXNet/NN.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN.hs
@@ -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
diff --git a/src/MXNet/NN/Callback.hs b/src/MXNet/NN/Callback.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/Callback.hs
@@ -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
diff --git a/src/MXNet/NN/DataIter/Class.hs b/src/MXNet/NN/DataIter/Class.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/DataIter/Class.hs
@@ -0,0 +1,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)
diff --git a/src/MXNet/NN/DataIter/Vec.hs b/src/MXNet/NN/DataIter/Vec.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/DataIter/Vec.hs
@@ -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
diff --git a/src/MXNet/NN/EvalMetric.hs b/src/MXNet/NN/EvalMetric.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/EvalMetric.hs
@@ -0,0 +1,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 :*
diff --git a/src/MXNet/NN/Initializer.hs b/src/MXNet/NN/Initializer.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/Initializer.hs
@@ -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"
diff --git a/src/MXNet/NN/Layer.hs b/src/MXNet/NN/Layer.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/Layer.hs
@@ -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
+
diff --git a/src/MXNet/NN/LrScheduler.hs b/src/MXNet/NN/LrScheduler.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/LrScheduler.hs
@@ -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)
diff --git a/src/MXNet/NN/NDArray.hs b/src/MXNet/NN/NDArray.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/NDArray.hs
@@ -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
diff --git a/src/MXNet/NN/Optimizer.hs b/src/MXNet/NN/Optimizer.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/Optimizer.hs
@@ -0,0 +1,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)
diff --git a/src/MXNet/NN/Types.hs b/src/MXNet/NN/Types.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/Types.hs
@@ -0,0 +1,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
diff --git a/src/MXNet/NN/Utils.hs b/src/MXNet/NN/Utils.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/Utils.hs
@@ -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)
diff --git a/src/MXNet/NN/Utils/GraphViz.hs b/src/MXNet/NN/Utils/GraphViz.hs
new file mode 100644
--- /dev/null
+++ b/src/MXNet/NN/Utils/GraphViz.hs
@@ -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