extensible-neural.lhs

% Extensible Neural Networks with Backprop
% Justin Le

This write-up is a follow-up to the *MNIST* tutorial
([rendered][mnist-rendered] here, and [literate haskell][mnist-lhs] here).
This write-up itself is available as a [literate haskell file][lhs], and
also [rendered as a pdf][rendered].

[mnist-rendered]: https://github.com/mstksg/backprop/blob/master/renders/backprop-mnist.pdf
[mnist-lhs]: https://github.com/mstksg/backprop/blob/master/samples/backprop-mnist.lhs
[rendered]: https://github.com/mstksg/backprop/blob/master/renders/extensible-neural.pdf
[lhs]: https://github.com/mstksg/backprop/blob/master/samples/extensible-neural.lhs

The (extra) packages involved are:

*   hmatrix
*   lens
*   mnist-idx
*   mwc-random
*   one-liner-instances
*   singletons
*   split

> {-# LANGUAGE BangPatterns        #-}
> {-# LANGUAGE DataKinds           #-}
> {-# LANGUAGE DeriveGeneric       #-}
> {-# LANGUAGE FlexibleContexts    #-}
> {-# LANGUAGE GADTs               #-}
> {-# LANGUAGE InstanceSigs        #-}
> {-# LANGUAGE LambdaCase          #-}
> {-# LANGUAGE RankNTypes          #-}
> {-# LANGUAGE ScopedTypeVariables #-}
> {-# LANGUAGE TemplateHaskell     #-}
> {-# LANGUAGE TypeApplications    #-}
> {-# LANGUAGE TypeInType          #-}
> {-# LANGUAGE TypeOperators       #-}
> {-# LANGUAGE ViewPatterns        #-}
> {-# OPTIONS_GHC -Wno-orphans     #-}
>
> import           Control.DeepSeq
> import           Control.Exception
> import           Control.Monad
> import           Control.Monad.IO.Class
> import           Control.Monad.Primitive
> import           Control.Monad.Trans.Maybe
> import           Control.Monad.Trans.State
> import           Data.Bitraversable
> import           Data.Foldable
> import           Data.IDX
> import           Data.Kind
> import           Data.List.Split
> import           Data.Singletons
> import           Data.Singletons.Prelude
> import           Data.Singletons.TypeLits
> import           Data.Time.Clock
> import           Data.Traversable
> import           Data.Tuple
> import           GHC.Generics                    (Generic)
> 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
============

The *[backprop][hackage]* library lets us manipulate our values in a
natural way.  We write the function to compute our result, and the library
then automatically finds the *gradient* of that function, which we can use
for gradient descent.

[hackage]: http://hackage.haskell.org/package/backprop

In the last post, we looked at using a fixed-structure neural network.
However, in [this blog series][blog], I discuss a system of extensible
neural networks that can be chained and composed.

[blog]: https://blog.jle.im/entries/series/+practical-dependent-types-in-haskell.html

One issue, however, in naively translating the implementations, is that we
normally run the network by pattern matching on each layer.  However, we
cannot directly pattern match on `BVar`s.

We *could* get around it by being smart with prisms and `^^?`, to extract a
"Maybe BVar".  However, we can do better!  This is because the *shape* of a
`Net i hs o` is known already at compile-time, so there is no need for
runtime checks like prisms and `^^?`.

Instead, we can just directly use lenses, since we know *exactly* what
constructor will be present!  We can use singletons to determine which
constructor is present, and so always just directly use lenses without any
runtime nondeterminism.

Types
=====

First, our types:

> data Layer i o =
>     Layer { _lWeights :: !(L o i)
>           , _lBiases  :: !(R o)
>           }
>   deriving (Show, Generic)
>
> instance NFData (Layer i o)
> makeLenses ''Layer
>
> data Net :: Nat -> [Nat] -> Nat -> Type where
>     NO   :: !(Layer i o) -> Net i '[] o
>     (:~) :: !(Layer i h) -> !(Net h hs o) -> Net i (h ': hs) o

Unfortunately, we can't automatically generate lenses for GADTs, so we have
to make them by hand.[^poly]

[^poly]: We write them originally as a polymorphic lens family to help us
with type safety via paraemtric polymorphism.

> _NO :: Lens (Net i '[] o) (Net i' '[] o')
>             (Layer i o  ) (Layer i' o'  )
> _NO f (NO l) = NO <$> f l
>
> _NIL :: Lens (Net i (h ': hs) o) (Net i' (h ': hs) o)
>              (Layer i h        ) (Layer i' h        )
> _NIL f (l :~ n) = (:~ n) <$> f l
>
> _NIN :: Lens (Net i (h ': hs) o) (Net i (h ': hs') o')
>              (Net h hs        o) (Net h hs'        o')
> _NIN f (l :~ n) = (l :~) <$> f n

You can read `_NO` as:

```haskell
_NO :: Lens' (Net i '[] o) (Layer i o)
```

A lens into a single-layer network, and

```haskell
_NIL :: Lens' (Net i (h ': hs) o) (Layer i h )
_NIN :: Lens' (Net i (h ': hs) o) (Net h hs o)
```

Lenses into a multiple-layer network, getting the first layer and the tail
of the network.

If we pattern match on `Sing hs`, we can always determine exactly which
lenses we can use, and so never fumble around with prisms or
nondeterminism.

Running the network
===================

Here's the meat of process, then: specifying how to run the network.  We
re-use our `BVar`-based combinators defined in the last write-up:

> 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 #-}

For `runNetwork`, we pattern match on `hs` using singletons, so we always
know exactly what type of network we have:

> runNetwork
>     :: (KnownNat i, KnownNat o, Reifies s W)
>     => BVar s (Net i hs o)
>     -> Sing hs
>     -> BVar s (R i)
>     -> BVar s (R o)
> runNetwork n = \case
>     SNil          -> softMax . runLayer (n ^^. _NO)
>     SCons SNat hs -> withSingI hs $
>                      runNetwork (n ^^. _NIN)  hs
>                    . logistic
>                    . runLayer (n ^^. _NIL)
> {-# INLINE runNetwork #-}

The rest of it is the same as before.

> netErr
>     :: (KnownNat i, KnownNat o, SingI hs, Reifies s W)
>     => R i
>     -> R o
>     -> BVar s (Net i hs o)
>     -> BVar s Double
> netErr x targ n = crossEntropy targ (runNetwork n sing (constVar x))
> {-# INLINE netErr #-}
>
> trainStep
>     :: forall i hs o. (KnownNat i, KnownNat o, SingI hs)
>     => Double             -- ^ learning rate
>     -> R i                -- ^ input
>     -> R o                -- ^ target
>     -> Net i hs o         -- ^ initial network
>     -> Net i hs o
> trainStep r !x !targ !n = n - realToFrac r * gradBP (netErr x targ) n
> {-# INLINE trainStep #-}
>
> trainList
>     :: (KnownNat i, SingI hs, KnownNat o)
>     => Double             -- ^ learning rate
>     -> [(R i, R o)]       -- ^ input and target pairs
>     -> Net i hs o         -- ^ initial network
>     -> Net i hs o
> trainList r = flip $ foldl' (\n (x,y) -> trainStep r x y n)
> {-# INLINE trainList #-}
>
> testNet
>     :: forall i hs o. (KnownNat i, KnownNat o, SingI hs)
>     => [(R i, R o)]
>     -> Net i hs 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 (\n' -> runNetwork n' sing (constVar x)) n

And that's it!

Running
=======

Everything here is the same as before, except now we can dynamically pick
the network size.  Here we pick `'[300,100]` for the hidden layer sizes.

> 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 @(Net 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

Looking Forward
===============

One common thing people might do is want to be able to mix different types
of layers.  This could also be easily encoded as different constructors in
`Layer`, and so `runLayer` will now be different depending on what
constructor is present.

In this case, we can either:

1.  Have a different indexed type for layers, so that we can always know
    exactly what layer is involved, so we don't have to runtime pattern
    match:

    ```haskell
    data LayerType = FullyConnected | Convolutional

    data Layer :: LayerType -> Nat -> Nat -> Type where
        LayerFC :: .... -> Layer 'FullyConnected i o
        LayerC  :: .... -> Layer 'Convolutional  i o
    ```

    We would then have `runLayer` take `Sing (t :: LayerType)`, so we can
    again use `^^.` and directly pattern match.

2.  Use a typeclass-based approach, so users can add their own layer types.
    In this situation, layer types would all be different types, and
    running them would be a typeclass method that would give our
    `BVar s (Layer i o) -> BVar s (R i) -> BVar s (R o)` operation as a
    typeclass method.

    ```haskell
    class Layer (l :: Nat -> Nat -> Type) where
        runLayer
            :: forall s. Reifies s W
            => BVar s (l i o)
            -> BVar s (R i)
            -> BVar s (R o)
    ```

In all cases, it shouldn't be much more cognitive overhead to use
*backprop* to build your neural network framework!

And, remember that `evalBP` (directly running the function) introduces
virtually zero overhead, so if you only provided `BVar` functions, you
could easily get the original non-`BVar` functions with `evalBP` without
any loss.

What now?
---------

Ready to start?  Check out the docs for the [Numeric.Backprop][] module for
the full technical specs, and find more examples and updates at the [github
repo][repo]!

[Numeric.Backprop]: http://hackage.haskell.org/package/backprop/docs/Numeric-Backprop.html
[repo]: https://github.com/mstksg/backprop

Internals
=========

That's it for the post!  Now for the internal plumbing :)

> 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)

HMatrix Operations
------------------

> 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) )
>
> 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)
>   )
>
> konst'
>     :: (KnownNat n, Reifies s W)
>     => BVar s Double
>     -> BVar s (R n)
> konst' = liftOp1 . op1 $ \c -> (konst c, HM.sumElements . extract)
>
> sumElements'
>     :: (KnownNat n, Reifies s W)
>     => BVar s (R n)
>     -> BVar s Double
> sumElements' = liftOp1 . op1 $ \x -> (HM.sumElements (extract x), konst)
>
> 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 #-}
>
> crossEntropy
>     :: (KnownNat n, Reifies s W)
>     => R n
>     -> BVar s (R n)
>     -> BVar s Double
> crossEntropy targ res = -(log res <.>! constVar targ)
> {-# INLINE crossEntropy #-}
>
> logistic :: Floating a => a -> a
> logistic x = 1 / (1 + exp (-x))
> {-# INLINE logistic #-}

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 o) => Fractional (Layer i o) where
>     (/)          = gDivide
>     recip        = gRecip
>     fromRational = gFromRational
>
> instance (KnownNat i, KnownNat o) => Backprop (Layer i o)
>
>
> liftNet0
>     :: forall i hs o. (KnownNat i, KnownNat o)
>     => (forall m n. (KnownNat m, KnownNat n) => Layer m n)
>     -> Sing hs
>     -> Net i hs o
> liftNet0 x = go
>   where
>     go :: forall w ws. KnownNat w => Sing ws -> Net w ws o
>     go = \case
>       SNil          -> NO x
>       SCons SNat hs -> x :~ go hs
>
> liftNet1
>     :: forall i hs o. (KnownNat i, KnownNat o)
>     => (forall m n. (KnownNat m, KnownNat n)
>           => Layer m n
>           -> Layer m n
>        )
>     -> Sing hs
>     -> Net i hs o
>     -> Net i hs o
> liftNet1 f = go
>   where
>     go  :: forall w ws. KnownNat w
>         => Sing ws
>         -> Net w ws o
>         -> Net w ws o
>     go = \case
>       SNil          -> \case
>         NO x -> NO (f x)
>       SCons SNat hs -> \case
>         x :~ xs -> f x :~ go hs xs
>
> liftNet2
>     :: forall i hs o. (KnownNat i, KnownNat o)
>     => (forall m n. (KnownNat m, KnownNat n)
>           => Layer m n
>           -> Layer m n
>           -> Layer m n
>        )
>     -> Sing hs
>     -> Net i hs o
>     -> Net i hs o
>     -> Net i hs o
> liftNet2 f = go
>   where
>     go  :: forall w ws. KnownNat w
>         => Sing ws
>         -> Net w ws o
>         -> Net w ws o
>         -> Net w ws o
>     go = \case
>       SNil          -> \case
>         NO x -> \case
>           NO y -> NO (f x y)
>       SCons SNat hs -> \case
>         x :~ xs -> \case
>           y :~ ys -> f x y :~ go hs xs ys
>
> instance ( KnownNat i
>          , KnownNat o
>          , SingI hs
>          )
>       => Num (Net i hs o) where
>     (+)           = liftNet2 (+) sing
>     (-)           = liftNet2 (-) sing
>     (*)           = liftNet2 (*) sing
>     negate        = liftNet1 negate sing
>     abs           = liftNet1 abs sing
>     signum        = liftNet1 signum sing
>     fromInteger x = liftNet0 (fromInteger x) sing
>
> instance ( KnownNat i
>          , KnownNat o
>          , SingI hs
>          )
>       => Fractional (Net i hs o) where
>     (/)            = liftNet2 (/) sing
>     recip          = liftNet1 negate sing
>     fromRational x = liftNet0 (fromRational x) sing
>
> instance (KnownNat i, KnownNat o, SingI hs) => Backprop (Net i hs o) where
>     zero = liftNet1 zero sing
>     add  = liftNet2 add sing
>     one  = liftNet1 one sing
>
> 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 o
>          , SingI hs
>          )
>       => MWC.Variate (Net i hs o) where
>     uniform :: forall m. PrimMonad m => MWC.Gen (PrimState m) -> m (Net i hs o)
>     uniform g = go sing
>       where
>         go :: forall w ws. KnownNat w => Sing ws -> m (Net w ws o)
>         go = \case
>           SNil          -> NO <$> MWC.uniform g
>           SCons SNat hs -> (:~) <$> MWC.uniform g <*> go hs
>     uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g
>
> instance NFData (Net i hs o) where
>     rnf = \case
>       NO l    -> rnf l
>       x :~ xs -> rnf x `seq` rnf xs
>
> 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