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/Dataset.hs b/examples/mnist/Dataset.hs new file mode 100644 --- /dev/null +++ b/examples/mnist/Dataset.hs @@ -0,0 +1,63 @@ +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE OverloadedLists #-} +module Dataset where + +import MXNet.Core.Base +import qualified MXNet.Core.Base.NDArray as A +import qualified MXNet.Core.Base.Internal.TH.NDArray as MXI +import Data.Function ((&)) +import Streaming +import Streaming.Prelude (Of(..)) +import qualified Streaming.Prelude as S +import Control.Monad.Trans.Resource (MonadResource(..)) +import qualified Data.Vector as NV +import qualified Data.Vector.Storable as SV + +import Parse + +type SymbolF = Symbol Float +type ArrayF = NDArray Float + +device :: Context +device = contextCPU + +type StreamProc a b m = Stream (Of a) m () -> Stream (Of b) m () + +mappedOf :: Monad m => (a -> m b) -> StreamProc a b m +-- mappedOf f = S.sequence . maps (first f) +mappedOf = S.mapM + +cImageToNDArray :: MonadIO m => StreamProc (Batched Image) ArrayF m +cImageToNDArray = mappedOf $ \dat -> liftIO $ do + let sz = size dat + makeNDArray [sz, 28, 28] device $ SV.concat $ NV.toList $ _batch dat + +cLabelToOnehotNDArray :: MonadIO m => StreamProc (Batched Label) ArrayF m +cLabelToOnehotNDArray = mappedOf $ \dat -> liftIO $ do + let sz = size dat + a <- array [sz] (NV.convert $ NV.map fromIntegral $ _batch dat) :: IO ArrayF + b <- MXI.one_hot (A.getHandle a) 10 (add @"on_value" 1.0 $ add @"off_value" 0.0 nil) + reshape (A.NDArray b) [sz, 10] + +cBatchN :: MonadIO m => Int -> StreamProc a (Batched a) m +cBatchN n = mapped toBatch . chunksOf n + where + toBatch seg = first (Batched . NV.fromList) <$> S.toList seg + +trainingData :: MonadResource m => Stream (Of (ArrayF, ArrayF)) m () +trainingData = S.zip + (sourceImages "examples/data/train-images-idx3-ubyte" & cBatchN 32 & cImageToNDArray ) + (sourceLabels "examples/data/train-labels-idx1-ubyte" & cBatchN 32 & cLabelToOnehotNDArray) + +testingData :: MonadResource m => Stream (Of (ArrayF, ArrayF)) m () +testingData = S.zip + (sourceImages "examples/data/t10k-images-idx3-ubyte" & cBatchN 1 & cImageToNDArray ) + (sourceLabels "examples/data/t10k-labels-idx1-ubyte" & cBatchN 1 & cLabelToOnehotNDArray) + +newtype Batched a = Batched { _batch :: NV.Vector a } + +size :: Batched a -> Int +size (Batched b) = NV.length b + 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,57 @@ +module Parse where + +import Streaming +import Data.Attoparsec.ByteString as AP +import Data.Attoparsec.Binary as AP +import Data.Attoparsec.ByteString.Streaming as APS +import qualified Data.ByteString.Streaming as BSS +import qualified Data.ByteString.Internal as BS +import qualified Data.Vector.Storable as SV +import Control.Exception.Base +import Control.Monad.Trans.Resource (MonadResource(..), MonadThrow(..)) +import Data.Typeable + +type Image = SV.Vector Float +type Label = Int + +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 + +sourceImages :: MonadResource m => FilePath -> Stream (Of Image) m () +sourceImages fp = do + (result, rest)<- lift $ APS.parse header (BSS.readFile fp) + case result of + Left (HeaderImg _ w h) -> void $ APS.parsed (image w h) rest + _ -> throwM NotImageFile + +sourceLabels :: MonadResource m => FilePath -> Stream (Of Label) m () +sourceLabels fp = do + (result, rest)<- lift $ APS.parse header (BSS.readFile fp) + case result of + Left (HeaderLbl _) -> void $ APS.parsed label rest + _ -> throwM NotImageFile + +data Exc = NotImageFile | NotLabelFile + deriving (Show, Typeable) +instance Exception Exc diff --git a/examples/mnist/mnist.hs b/examples/mnist/mnist.hs new file mode 100644 --- /dev/null +++ b/examples/mnist/mnist.hs @@ -0,0 +1,74 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +module Main where + +import MXNet.Core.Base +import qualified MXNet.Core.Base.NDArray as A +import qualified MXNet.Core.Base.Internal.TH.NDArray as A +import qualified Data.HashMap.Strict as M +import Control.Monad (forM_) +import qualified Streaming.Prelude as SR +import qualified Data.Vector.Storable as SV +import Data.List (intersperse) +import Control.Monad.IO.Class +import Control.Monad.Trans.Resource +import MXNet.NN +import Dataset + +neural :: IO SymbolF +neural = do + x <- variable "x" :: IO SymbolF + y <- variable "y" :: IO SymbolF + w1 <- variable "w1" :: IO SymbolF + b1 <- variable "b1" :: IO SymbolF + v1 <- fullyConnected x w1 b1 128 + a1 <- activation v1 "relu" + w2 <- variable "w2" :: IO SymbolF + b2 <- variable "b2" :: IO SymbolF + v2 <- fullyConnected a1 w2 b2 10 + a2 <- softmaxOutput v2 y + return a2 + +range :: Int -> [Int] +range = enumFromTo 1 + +default_initializer :: DType a => [Int] -> IO (NDArray a) +default_initializer shape = A.NDArray <$> A.random_normal (add @"loc" 0 $ add @"scale" 1 $ add @"shape" formatedShape nil) + where + formatedShape = concat $ ["("] ++ intersperse "," (map show shape) ++ [")"] + +optimizer :: DType a => NDArray a -> NDArray a -> IO (NDArray a) +optimizer v g = A.NDArray <$> (A.sgd_update (A.getHandle v) (A.getHandle g) 0.01 nil) + +main :: IO () +main = do + -- call mxListAllOpNames can ensure the MXNet itself is properly initialized + -- i.e. MXNet operators are registered in the NNVM + _ <- mxListAllOpNames + net <- neural + params <- initialize net $ Config { + _cfg_placeholders = M.singleton "x" [32,28,28], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer + } + result <- runResourceT $ train params contextCPU $ do + liftIO $ putStrLn $ "[Train] " + forM_ (range 5) $ \ind -> do + liftIO $ putStrLn $ "iteration " ++ show ind + SR.mapM_ (\(x, y) -> fit optimizer net $ M.fromList [("x", x), ("y", y)]) trainingData + liftIO $ putStrLn $ "[Test] " + SR.toList_ $ flip SR.mapM testingData $ \(x, y) -> do + [y'] <- forwardOnly net (M.fromList [("x", Just x), ("y", Nothing)]) + ind1 <- liftIO $ argmax y >>= items + ind2 <- liftIO $ argmax y' >>= items + return (ind1, ind2) + let (ls,ps) = unzip result + ls_unbatched = mconcat ls + ps_unbatched = mconcat ps + total = SV.length ls_unbatched + correct = SV.length $ SV.filter id $ SV.zipWith (==) ls_unbatched ps_unbatched + putStrLn $ "Accuracy: " ++ show correct ++ "/" ++ show total + + where + argmax :: ArrayF -> IO ArrayF + argmax ys = A.NDArray <$> A.argmax (A.getHandle ys) (add @"axis" 1 nil) diff --git a/mxnet-nn.cabal b/mxnet-nn.cabal new file mode 100644 --- /dev/null +++ b/mxnet-nn.cabal @@ -0,0 +1,50 @@ +name: mxnet-nn +version: 0.0.1 +synopsis: Train a neural network with MXNet in Haskell. +description: High level APIs to rain a neural network with MXNet in Haskell. +homepage: http://github.com/pierric/mxnet-haskell-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 + other-modules: + hs-source-dirs: src + ghc-options: -Wall + default-language: Haskell2010 + build-depends: base >= 4.7 && < 5.0 + , mxnet >= 0.2.0.0 + , unordered-containers >= 0.2.8 + , resourcet >= 1.1.8 + , vector >= 0.12 + , mtl >= 2.2 + , lens >= 4.12 + +Executable mnist + main-is: mnist.hs + other-modules: Parse Dataset + hs-source-dirs: examples/mnist + ghc-options: -Wall + default-language: Haskell2010 + build-depends: base >= 4.7 && < 5.0 + , mxnet >= 0.2.0.0 + , unordered-containers >= 0.2.8 + , attoparsec >= 0.13 + , attoparsec-binary >= 0.2 + , vector >= 0.12 + , bytestring >= 0.10 + , resourcet >= 1.1.8 + , exceptions >= 0.8.3 + , mmorph >= 1.0.9 + , mtl >= 2.2.0 + , streaming >= 0.1.4.5 + , streaming-utils >= 0.1.4.5 + , streaming-bytestring >= 0.1.4.5 + , ghc-prim + , mxnet-nn 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,180 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE RecordWildCards #-} +module MXNet.NN ( + Parameter(..), + Config(..), + Exc(..), + Initializer, + Optimizer, + TrainM, + train, + inferShape, + initialize, + fit, + forwardOnly +) where + +import MXNet.Core.Base hiding (bind, context) +import MXNet.Core.Base.Internal +import qualified MXNet.Core.Base.NDArray as A +import qualified MXNet.Core.Base.Symbol as S +import qualified MXNet.Core.Base.Executor as E +import qualified MXNet.Core.Types.Internal as MXI +import qualified Data.HashMap.Strict as M +import Data.Typeable +import qualified Control.Monad.State as ST +import Data.Maybe (isJust, fromJust) +import Control.Monad (when) +import Control.Monad.IO.Class (MonadIO, liftIO) +import Control.Monad.Trans.Resource (MonadThrow(..)) +import Control.Exception.Base (Exception) +import Control.Lens (traverseOf, _1) + +-- | A parameter is two 'NDArray' to back a 'Symbol' +data Parameter a = Parameter { _param_in :: NDArray a, _param_grad :: NDArray a } + deriving Show + +-- | TrainM is a 'StateT' monad, where the state is all the 'Parameters' and a 'Context' +type TrainM a m = ST.StateT (M.HashMap String (Parameter a), Context) m + +-- | Initializer is about how to create a NDArray from a given shape. +-- +-- Usually, it can be a wrapper of MXNet operators, such as @random_uniform@, @random_normal@, +-- @random_gamma@, etc.. +type Initializer a = [Int] -> IO (NDArray a) +type Optimizer a = NDArray a -> NDArray a -> IO (NDArray a) + +-- | Execute the 'TrainM' monad +train :: (DType a, Monad m) => M.HashMap String (Parameter a) -> Context -> TrainM a m r -> m r +train param context = flip ST.evalStateT (param, context) + +-- | infer the shapes of all the symbols in a symbolic neural network +inferShape :: DType a => Symbol a -> M.HashMap String (NDArray a) -> IO (M.HashMap String [Int]) +inferShape sym known = do + let (names, vals) = unzip $ M.toList known + shapes <- mapM ndshape vals + let arg_ind = scanl (+) 0 $ map fst shapes + arg_shp = concat $ map snd shapes + (inp_shp, _, _) <- mxSymbolInferShape (S.getHandle sym) names arg_ind arg_shp + inps <- listInputs sym + return $ M.fromList $ zip inps inp_shp + +-- | 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_placeholders :: M.HashMap String [Int], + _cfg_initializers :: M.HashMap String (Initializer a), + _cfg_default_initializer :: Initializer a +} + +-- | initialize all parameters +initialize :: DType a => Symbol a -> Config a -> IO (M.HashMap String (Parameter a)) +initialize sym config = do + let spec1 = M.difference (_cfg_placeholders config) (_cfg_initializers config) + spec2 = _cfg_initializers config + dinit = _cfg_default_initializer config + placeholder <- mapM zeros spec1 + inp_with_shp <- inferShape sym placeholder + M.traverseWithKey (init_with_random_normal placeholder spec2 dinit) inp_with_shp + where + init_with_random_normal placeholder spec2 dinit inp shp = do + case M.lookup inp placeholder of + Just in_arg -> return $ Parameter in_arg (A.NDArray MXI.nullNDArrayHandle) + Nothing -> do + arg_in <- case M.lookup inp spec2 of + Just cinit -> cinit shp + Nothing -> dinit shp + arg_gr <- zeros shp + return $ Parameter arg_in arg_gr + +-- | bind the symbolic network with actual parameters +bind :: DType a => Symbol a -> M.HashMap String (Parameter a) -> Context -> Bool -> IO (Executor a) +bind net args Context{..} train_ = do + names <- listInputs net + exec_handle <- checked $ mxExecutorBind (S.getHandle net) deviceType deviceId + (fromIntegral (M.size args)) + -- the parameters to bind should be arranged in the same order as the names + (map (A.getHandle . _param_in) $ map (args M.!) names) + (if train_ + then map (A.getHandle . _param_grad) $ map (args M.!) names + else replicate (M.size args) MXI.nullNDArrayHandle) + (replicate (M.size args) 1) + 0 [] + + makeExecutor exec_handle + +-- | single step train. Must provide all the placeholders. +fit :: (DType a, MonadIO m, MonadThrow m) => Optimizer a -> Symbol a -> M.HashMap String (NDArray a) -> TrainM a m () +fit opt net datAndLbl = do + shps <- liftIO $ inferShape net datAndLbl + modifyT . traverseOf _1 $ M.traverseWithKey $ \k p -> do + let ishp = shps M.! k + case M.lookup k datAndLbl of + Just a -> return $ p {_param_in = a} + Nothing -> do + (_, pshp1) <- liftIO $ ndshape (_param_in p) + (_, pshp2) <- liftIO $ ndshape (_param_grad p) + when (ishp /= pshp1 || ishp /= pshp2) (throwM $ MismatchedShape k) + return p + (params, context) <- ST.get + liftIO $ do + exec <- bind net params context True + checked $ mxExecutorForward (E.getHandle exec) 1 + backward exec + modifyT . traverseOf _1 $ M.traverseWithKey $ \ k v -> do + if (not $ M.member k datAndLbl) + then do new_in <- liftIO $ opt (_param_in v) (_param_grad v) + return $ v {_param_in = new_in} + else return v + +-- | 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) => Symbol a -> M.HashMap String (Maybe (NDArray a)) -> TrainM a m [NDArray a] +forwardOnly net dat = do + shps <- liftIO $ inferShape net (M.map fromJust $ M.filter isJust dat) + modifyT . traverseOf _1 $ M.traverseWithKey $ \k p -> do + let ishp = shps M.! k + case M.lookup k dat of + Just (Just a) -> + return $ p {_param_in = a} + Just Nothing -> do + dummy <- liftIO $ zeros ishp + return $ p {_param_in = dummy} + Nothing -> do + (_, pshp) <- liftIO $ ndshape (_param_in p) + when (ishp /= pshp) (throwM $ MismatchedShape k) + return p + (params, context) <- ST.get + liftIO $ do + exec <- bind net params context False + checked $ mxExecutorForward (E.getHandle exec) 0 + getOutputs exec + +-- | Possible exception in 'TrainM' +data Exc = MismatchedShape String + deriving (Show, Typeable) +instance Exception Exc + +-- | 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 +