% Learning MNIST with Neural Networks with backprop library % Justin Le The *backprop* library performs back-propagation over a *hetereogeneous* system of relationships. back-propagation is done automatically (as reverse-mode automatic differentiation), and you work with your values as if you were writing normal functions with them, with the help of [lens][]. [ad]: http://hackage.haskell.org/package/ad [lens]: http://hackage.haskell.org/package/lens Repository source is [on github][repo], and docs are [on hackage][hackage]. [repo]: https://github.com/mstksg/backprop [hackage]: http://hackage.haskell.org/package/backprop If you're reading this as a literate haskell file, you should know that a [rendered pdf version is available on github.][rendered]. If you are reading this as a pdf file, you should know that a [literate haskell version that you can run][lhs] is also available on github! [rendered]: https://github.com/mstksg/backprop/blob/master/renders/backprop-mnist.pdf [lhs]: https://github.com/mstksg/backprop/blob/master/samples/backprop-mnist.lhs The (extra) packages involved are: * hmatrix * lens * mnist-idx * mwc-random * one-liner-instances * split > {-# LANGUAGE BangPatterns #-} > {-# LANGUAGE DataKinds #-} > {-# LANGUAGE DeriveGeneric #-} > {-# LANGUAGE FlexibleContexts #-} > {-# LANGUAGE GADTs #-} > {-# LANGUAGE LambdaCase #-} > {-# LANGUAGE ScopedTypeVariables #-} > {-# LANGUAGE TemplateHaskell #-} > {-# LANGUAGE TupleSections #-} > {-# LANGUAGE TypeApplications #-} > {-# LANGUAGE ViewPatterns #-} > {-# OPTIONS_GHC -Wno-incomplete-patterns #-} > {-# OPTIONS_GHC -Wno-orphans #-} > {-# OPTIONS_GHC -Wno-unused-top-binds #-} > > import Control.DeepSeq > import Control.Exception > import Control.Monad > import Control.Monad.IO.Class > import Control.Monad.Trans.Maybe > import Control.Monad.Trans.State > import Data.Bitraversable > import Data.Foldable > import Data.IDX > import Data.List.Split > import Data.Time.Clock > import Data.Traversable > import Data.Tuple > import GHC.Generics (Generic) > import GHC.TypeLits > import Lens.Micro > import Lens.Micro.TH > import Numeric.Backprop > import Numeric.Backprop.Class > import Numeric.LinearAlgebra.Static > import Numeric.OneLiner > import Text.Printf > import qualified Data.Vector as V > import qualified Data.Vector.Generic as VG > import qualified Data.Vector.Unboxed as VU > import qualified Numeric.LinearAlgebra as HM > import qualified System.Random.MWC as MWC > import qualified System.Random.MWC.Distributions as MWC Introduction ============ In this walkthrough, we'll be building a classifier for the *[MNIST][]* data set. This is meant to mirror the [Tensorflow Tutorial][tf-intro] for beginners. [tf-intro]: https://www.tensorflow.org/versions/r1.2/get_started/mnist/beginners Essentially, we use a two-layer artificial neural network -- or a series of matrix multiplications, differentiable function applications, and vector additions. We feed our input image to the ANN and then try to get a label from it. Training an ANN is a matter of finding the right matrices to multiply by, and the right vectors to add. To do that, we train our network by treating our network's accuracy as a function Network -> Error. If we can find the gradient of the input network with respect to the error, we can perform [gradient descent][], and slowly make our network better and better. [gradient descent]: https://en.wikipedia.org/wiki/Gradient_descent Finding the gradient is usually complicated, but *backprop* makes it simpler: 1. Write a function to compute the error from the network 2. That's it! Hooray! Once you do that, the library finds the gradient function *automatically*, without any further intervention! Types ===== For the most part, we're going to be using the great *[hmatrix][]* library and its vector and matrix types. It offers a type L m n for $m \times n$ matrices, and a type R n for an $n$ vector. [hmatrix]: http://hackage.haskell.org/package/hmatrix First things first: let's define our neural networks as simple containers of parameters (weight matrices and bias vectors). First, a type for layers: > data Layer i o = > Layer { _lWeights :: !(L o i) > , _lBiases :: !(R o) > } > deriving (Show, Generic) > > instance NFData (Layer i o) > makeLenses ''Layer And a type for a simple feed-forward network with two hidden layers: > data Network i h1 h2 o = > Net { _nLayer1 :: !(Layer i h1) > , _nLayer2 :: !(Layer h1 h2) > , _nLayer3 :: !(Layer h2 o) > } > deriving (Show, Generic) > > instance NFData (Network i h1 h2 o) > makeLenses ''Network These are pretty straightforward container types...pretty much exactly the type you'd make to represent these networks! Note that, following true Haskell form, we separate out logic from data. This should be all we need. Instances --------- Things are much simplier if we had Num and Fractional instances for everything, so let's just go ahead and define that now, as well. Just a little bit of boilerplate, made easier using *[one-liner-instances][]* to auto-derive instances using Generics. [one-liner-instances]: http://hackage.haskell.org/package/one-liner-instances > instance (KnownNat i, KnownNat o) => Num (Layer i o) where > (+) = gPlus > (-) = gMinus > (*) = gTimes > negate = gNegate > abs = gAbs > signum = gSignum > fromInteger = gFromInteger > > instance ( KnownNat i > , KnownNat h1 > , KnownNat h2 > , KnownNat o > ) => Num (Network i h1 h2 o) where > (+) = gPlus > (-) = gMinus > (*) = gTimes > negate = gNegate > abs = gAbs > signum = gSignum > fromInteger = gFromInteger > > instance (KnownNat i, KnownNat o) => Fractional (Layer i o) where > (/) = gDivide > recip = gRecip > fromRational = gFromRational > > instance ( KnownNat i > , KnownNat h1 > , KnownNat h2 > , KnownNat o > ) => Fractional (Network i h1 h2 o) where > (/) = gDivide > recip = gRecip > fromRational = gFromRational KnownNat comes from *base*; it's a typeclass that *hmatrix* uses to refer to the numbers in its type and use it to go about its normal hmatrixy business. Now we need instances of Backprop for our types in order to use them for automatic differentiation. Luckily, these can be generated automatically using GHC Generics: > instance (KnownNat i, KnownNat o) => Backprop (Layer i o) > instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Backprop (Network i h1 h2 o) Ops === Now, *backprop* does require *primitive* differentiable operations on our relevant types to be defined. *backprop* uses these primitive operations to tie everything together. Ideally we'd import these from a library that implements these for you, and the end-user never has to make these primitives. But in this case, I'm going to put the definitions here to show that there isn't any magic going on. If you're curious, refer to [documentation for Op][opdoc] for more details on how Op is implemented and how this works. [opdoc]: http://hackage.haskell.org/package/backprop/docs/Numeric-Backprop-Op.html First, matrix-vector multiplication primitive, giving an explicit gradient function. > infixr 8 #>! > (#>!) > :: (KnownNat m, KnownNat n, Reifies s W) > => BVar s (L m n) > -> BVar s (R n) > -> BVar s (R m) > (#>!) = liftOp2 . op2 $\m v -> > ( m #> v, \g -> (g outer v, tr m #> g) ) Dot products would be nice too. > infixr 8 <.>! > (<.>!) > :: (KnownNat n, Reifies s W) > => BVar s (R n) > -> BVar s (R n) > -> BVar s Double > (<.>!) = liftOp2 . op2$ \x y -> > ( x <.> y, \g -> (konst g * y, x * konst g) > ) Also a function to fill a vector with the same element: > konst' > :: (KnownNat n, Reifies s W) > => BVar s Double > -> BVar s (R n) > konst' = liftOp1 . op1 $\c -> (konst c, HM.sumElements . extract) Finally, an operation to sum all of the items in the vector. > sumElements' > :: (KnownNat n, Reifies s W) > => BVar s (R n) > -> BVar s Double > sumElements' = liftOp1 . op1$ \x -> (HM.sumElements (extract x), konst) Again, these are not intended to be used by end-users of *backprop*, but rather are meant to be provided by libraries as primitive operations for users of the library to use. Running our Network =================== Now that we have our primitives in place, let's actually write a function to run our network! And, once we do this, we automatically also have functions to back-propagate our network! Normally, to write this function, we'd write: > runLayerNormal > :: (KnownNat i, KnownNat o) > => Layer i o > -> R i > -> R o > runLayerNormal l x = (l ^. lWeights) #> x + (l ^. lBiases) > {-# INLINE runLayerNormal #-} Using the lWeights and lBiases lenses to access the weights and biases of our layer. However, we can translate this to *backprop* by operating on BVars instead of the type directly, and using our backprop-aware #>!: > runLayer > :: (KnownNat i, KnownNat o, Reifies s W) > => BVar s (Layer i o) > -> BVar s (R i) > -> BVar s (R o) > runLayer l x = (l ^^. lWeights) #>! x + (l ^^. lBiases) > {-# INLINE runLayer #-} ^. lets to access data within a value using a lens, and ^^. lets you access data within a BVar using a lens: haskell (^.) :: a -> Lens' a b -> b (^^.) :: BVar s a -> Lens' a b -> BVar s b  (There is also ^^?, which can use a Prism or Traversal to extract a target that might not exist, ^^.., which uses a Traversal to extract all targets, and .~~, which uses a Lens to update a value inside BVar) Now runLayer is a function on two inputs that can be backpropagated, automatically! We can find its gradient given any input, and also run it to get our expected output as well. Before writing our final network runner, we need a function to compute the "softmax" of our output vector. Writing it normally would look like: > softMaxNormal :: KnownNat n => R n -> R n > softMaxNormal x = konst (1 / HM.sumElements (extract expx)) * expx > where > expx = exp x > {-# INLINE softMaxNormal #-} But we can make the mechanical shift to the backpropagatable version: > softMax :: (KnownNat n, Reifies s W) => BVar s (R n) -> BVar s (R n) > softMax x = konst' (1 / sumElements' expx) * expx > where > expx = exp x > {-# INLINE softMax #-} We also need the [logistic function][], which is our activation function between layer outputs. Because BVars have a Floating instance, we can just write it using typeclass functions. [logistic function]: https://en.wikipedia.org/wiki/Logistic_function > logistic :: Floating a => a -> a > logistic x = 1 / (1 + exp (-x)) > {-# INLINE logistic #-} With those in hand, let's compare how we would normally write a function to run our network: > runNetNormal > :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) > => Network i h1 h2 o > -> R i > -> R o > runNetNormal n = softMaxNormal > . runLayerNormal (n ^. nLayer3) > . logistic > . runLayerNormal (n ^. nLayer2) > . logistic > . runLayerNormal (n ^. nLayer1) > {-# INLINE runNetNormal #-} Basic function composition, neat. We use our lenses nLayer1, nLayer2, and nLayer3 to extract the first, second, and third layers from our network. Writing it in a way that backprop can use is also very similar: > runNetwork > :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) > => BVar s (Network i h1 h2 o) > -> R i > -> BVar s (R o) > runNetwork n = softMax > . runLayer (n ^^. nLayer3) > . logistic > . runLayer (n ^^. nLayer2) > . logistic > . runLayer (n ^^. nLayer1) > . constVar > {-# INLINE runNetwork #-} We use constVar on the input vector, because we don't care about its gradient and so treat it as a constant. And now here again we use ^^. (instead of ^.) to extract a value from our BVar of a Network, using a lens. Computing Errors ---------------- Now, training a neural network is about calculating its gradient with respect to some error function. The library calculatues the gradient for us -- we just need to tell it how to compute the error function. For classification problems, we usually use a [cross entropy][] error. Given a target vector, how does our neural network's output differ from what is expected? Lower numbers are better! [cross entropy]: https://en.wikipedia.org/wiki/Cross_entropy Again, let's look at a "normal" implementation, regular variables and no backprop: > crossEntropyNormal :: KnownNat n => R n -> R n -> Double > crossEntropyNormal targ res = -(log res <.> targ) > {-# INLINE crossEntropyNormal #-} And we can see that the backpropable version is pretty similar. We see constVar t, to introduce a BVar that is a constant value (that we don't care about the gradient of). > crossEntropy > :: (KnownNat n, Reifies s W) > => R n > -> BVar s (R n) > -> BVar s Double > crossEntropy targ res = -(log res <.>! constVar targ) > {-# INLINE crossEntropy #-} Our final "error function", then, is: > netErr > :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) > => R i > -> R o > -> BVar s (Network i h1 h2 o) > -> BVar s Double > netErr x targ n = crossEntropy targ (runNetwork n x) > {-# INLINE netErr #-} The Magic ========= The actual "magic" of the library happens with the functions to "run" the functions we defined earlier: haskell evalBP :: (forall s. Reifies s W => BVar s a -> BVar s b) -> a -> b gradBP :: (forall s. Reifies s W => BVar s a -> BVar s b) -> a -> a backprop :: (forall s. Reifies s W => BVar s a -> BVar s b) -> a -> (b, a)  evalBP "runs" the function like normal, gradBP computes the gradient of the function, and backprop computes both the result and the gradient. So, if we have a network net0, an input vector x, and a target vector t, we could compute its error using: haskell evalBP (netErr x targ) net0 :: Double  And we can calculate its *gradient* using: haskell gradBP (netErr x targ) net0 :: (Network i h1 h2 o, R i)  Pulling it all together ======================= Let's write a simple function to step our network in the direction opposite of the gradient to train our model: > trainStep > :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) > => Double -- ^ learning rate > -> R i -- ^ input > -> R o -- ^ target > -> Network i h1 h2 o -- ^ initial network > -> Network i h1 h2 o > trainStep r !x !targ !n = n - realToFrac r * gradBP (netErr x targ) n > {-# INLINE trainStep #-} Here's a convenient wrapper for training over all of the observations in a list: > trainList > :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) > => Double -- ^ learning rate > -> [(R i, R o)] -- ^ input and target pairs > -> Network i h1 h2 o -- ^ initial network > -> Network i h1 h2 o > trainList r = flip $foldl' (\n (x,y) -> trainStep r x y n) > {-# INLINE trainList #-} testNet will be a quick way to test our net by computing the percentage of correct guesses: (mostly using *hmatrix* stuff, so don't mind too much) > testNet > :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) > => [(R i, R o)] > -> Network i h1 h2 o > -> Double > testNet xs n = sum (map (uncurry test) xs) / fromIntegral (length xs) > where > test :: R i -> R o -> Double -- test if the max index is correct > test x (extract->t) > | HM.maxIndex t == HM.maxIndex (extract r) = 1 > | otherwise = 0 > where > r :: R o > r = evalBP (runNetwork x) n And now, a main loop! If you are following along at home, download the [mnist data set files][MNIST] and uncompress them into the folder data, and everything should work fine. [MNIST]: http://yann.lecun.com/exdb/mnist/ > main :: IO () > main = MWC.withSystemRandom$ \g -> do > Just train <- loadMNIST "data/train-images-idx3-ubyte" "data/train-labels-idx1-ubyte" > Just test <- loadMNIST "data/t10k-images-idx3-ubyte" "data/t10k-labels-idx1-ubyte" > putStrLn "Loaded data." > net0 <- MWC.uniformR @(Network 784 300 100 10) (-0.5, 0.5) g > flip evalStateT net0 . forM_ [1..] $\e -> do > train' <- liftIO . fmap V.toList$ MWC.uniformShuffle (V.fromList train) g > liftIO $printf "[Epoch %d]\n" (e :: Int) > > forM_ ([1..] zip chunksOf batch train')$ \(b, chnk) -> StateT $\n0 -> do > printf "(Batch %d)\n" (b :: Int) > > t0 <- getCurrentTime > n' <- evaluate . force$ trainList rate chnk n0 > t1 <- getCurrentTime > printf "Trained on %d points in %s.\n" batch (show (t1 diffUTCTime t0)) > > let trainScore = testNet chnk n' > testScore = testNet test n' > printf "Training error: %.2f%%\n" ((1 - trainScore) * 100) > printf "Validation error: %.2f%%\n" ((1 - testScore ) * 100) > > return ((), n') > where > rate = 0.02 > batch = 5000 Each iteration of the loop: 1. Shuffles the training set 2. Splits it into chunks of batch size 3. Uses trainList to train over the batch 4. Computes the score based on testNet based on the training set and the test set 5. Prints out the results And, that's really it! Performance ----------- Currently, benchmarks show that *running* the network has virtually zero overhead (~ 4%) over writing the running function directly. The actual gradient descent process (compute gradient, then descend) carries about 60% overhead over writing the gradients manually, but it is unclear how much of this is because of the library, and how much of it is just because of automatic differentation giving slightly less efficient matrix/vector multiplication operations. The [README][repo] has some more detailed benchmarks and statistics, if you want to get more detailed information. Main takeaways ============== Most of the actual heavy lifting/logic actually came from the *hmatrix* library itself. We just created simple types to wrap up our bare matrices. Basically, all that *backprop* did was give you an API to define *how to run* a neural net --- how to *run* a net based on a Network and R i input you were given. The goal of the library is to let you write down how to run things in as natural way as possible. And then, after things are run, we can just get the gradient and roll from there! Because the heavy lifting is done by the data types themselves, we can presumably plug in *any* type and any tensor/numerical backend, and reap the benefits of those libraries' optimizations and parallelizations. *Any* type can be backpropagated! :D What now? --------- Check out the docs for the [Numeric.Backprop][] module for a more detailed picture of what's going on, or find more examples at the [github repo][repo]! [Numeric.Backprop]: http://hackage.haskell.org/package/backprop/docs/Numeric-Backprop.html Also, check out follow-up writeup to this tutorial, expanding on using the library with more advanced extensible neural network types, like the ones described in [this blog post][blog]. Check out the [literate haskell here][neural-lhs], and the [rendered PDF here][neural-pdf]. [blog]: https://blog.jle.im/entries/series/+practical-dependent-types-in-haskell.html [neural-lhs]: https://github.com/mstksg/backprop/blob/master/samples/extensible-neural.lhs [neural-pdf]: https://github.com/mstksg/backprop/blob/master/renders/extensible-neural.pdf Boring stuff ============ Here is a small wrapper function over the [mnist-idx][] library loading the contents of the idx files into *hmatrix* vectors: [mnist-idx]: http://hackage.haskell.org/package/mnist-idx > loadMNIST > :: FilePath > -> FilePath > -> IO (Maybe [(R 784, R 10)]) > loadMNIST fpI fpL = runMaybeT $do > i <- MaybeT$ decodeIDXFile fpI > l <- MaybeT $decodeIDXLabelsFile fpL > d <- MaybeT . return$ labeledIntData l i > r <- MaybeT . return $for d (bitraverse mkImage mkLabel . swap) > liftIO . evaluate$ force r > where > mkImage :: VU.Vector Int -> Maybe (R 784) > mkImage = create . VG.convert . VG.map (\i -> fromIntegral i / 255) > mkLabel :: Int -> Maybe (R 10) > mkLabel n = create $HM.build 10 (\i -> if round i == n then 1 else 0) And here are instances to generating random vectors/matrices/layers/networks, used for the initialization step. > instance KnownNat n => MWC.Variate (R n) where > uniform g = randomVector <$> MWC.uniform g <*> pure Uniform > uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g > > instance (KnownNat m, KnownNat n) => MWC.Variate (L m n) where > uniform g = uniformSample <$> MWC.uniform g <*> pure 0 <*> pure 1 > uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g > > instance (KnownNat i, KnownNat o) => MWC.Variate (Layer i o) where > uniform g = Layer <$> MWC.uniform g <*> MWC.uniform g > uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g > > instance ( KnownNat i > , KnownNat h1 > , KnownNat h2 > , KnownNat o > ) > => MWC.Variate (Network i h1 h2 o) where > uniform g = Net <$> MWC.uniform g <*> MWC.uniform g <*> MWC.uniform g > uniformR (l, h) g = (\x -> x * (h - l) + l) <\$> MWC.uniform g Also, some orphan instances of Backprop for vector and matrix types. These are provided by the [hmatrix-backprop][] library normally: > instance Backprop (R n) where > zero = zeroNum > add = addNum > one = oneNum > > instance (KnownNat n, KnownNat m) => Backprop (L m n) where > zero = zeroNum > add = addNum > one = oneNum [hmatrix-backprop]: http://hackage.haskell.org/package/hmatrix-backprop