MNIST.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}

module Main where

import           Codec.Picture
import           Control.Category
import qualified Data.Array       as A
import           Data.MyPrelude
import           Data.Utils
import           Numeric.Neural
import           Pipes.GZip       (decompress)
import qualified Pipes.Prelude    as P
import           Prelude          hiding (id, (.))

main :: IO ()
main = flip evalRandT (mkStdGen 999999) $ do
    xs     <- getSamples [0 .. 999]
    m      <- modelR (whiten mnistModel $ fst <$> xs)
    liftIO $ printf "generation  learning rate  batch error\n\n"
    (a, g) <- runEffect $
            cachingBatchP getSamples 60000 20 2000 100
        >-> descentP m 1 (\g -> 0.4 * 100 / (100 + fromIntegral g))
        >-> reportTSP 1 report
        >-> consumeTSP check
    liftIO $ printf "\nreached accuracy of %f after %d generations\n" a g

  where

    getSamples xs = liftIO $ runSafeT $ P.toListM $ trainSamples >-> indicesP xs

    report ts = liftIO $ do
        let g = tsGeneration ts
        when (g `mod` 5 == 0) $ printf "   %7d       %8.6f   %10.8f\n" g (tsEta ts) (tsBatchError ts)

    check ts = do
        let g = tsGeneration ts
        if g `mod` 50 == 0
            then do
                a <- liftIO $ accuracy $ tsModel ts
                liftIO $ printf "\naccuracy %f\n\n" a
                return $ if a > 0.9 then Just (a, g) else Nothing
            else return Nothing

accuracy :: MNISTModel -> IO Double
accuracy m = runSafeT $ fromJust <$> classifierAccuracyP m testSamples

type Img = Image Pixel8

data Digit = Zero | One | Two | Three | Four | Five | Six | Seven | Eight | Nine
    deriving (Show, Read, Eq, Ord, Enum, Bounded)

type Sample = (Img, Digit)

trainImagesFile, trainLabelsFile, testImagesFile, testLabelsFile :: FilePath
trainImagesFile = "examples" </> "MNIST" </> "train-images-idx3-ubyte" <.> "gz"
trainLabelsFile = "examples" </> "MNIST" </> "train-labels-idx1-ubyte" <.> "gz"
testImagesFile  = "examples" </> "MNIST" </> "t10k-images-idx3-ubyte"  <.> "gz"
testLabelsFile  = "examples" </> "MNIST" </> "t10k-labels-idx1-ubyte"  <.> "gz"

bytes :: MonadSafe m => FilePath -> Producer Word8 m ()
bytes f = decompress (fromFile f) >-> toWord8

labels :: MonadSafe m => FilePath -> Producer Digit m ()
labels f = bytes f >-> P.drop 8 >-> P.map (toEnum . fromIntegral)

images :: MonadSafe m => FilePath -> Producer Img m ()
images f = bytes f >-> P.drop 16 >-> chunks (28 * 28) >-> P.map g

  where

    g xs = let a = A.listArray ((0, 0), (27, 27)) xs
           in  generateImage (\x y -> 255 - a A.! (y, x)) 28 28

trainSamples, testSamples :: MonadSafe m => Producer Sample m ()
trainSamples = P.zip (images trainImagesFile) (labels trainLabelsFile)
testSamples  = P.zip (images testImagesFile)  (labels testLabelsFile)

writeImg :: MonadIO m => FilePath -> Img -> m ()
writeImg f i = liftIO $ saveTiffImage (f <.> "tiff") (ImageY8 i)

type MNISTModel = Classifier (Matrix 28 28) 10 Img Digit

mnistModel :: MNISTModel
mnistModel = mkStdClassifier c i where

    c = tanhLayer' . (tanhLayer' :: Layer 784 30) . cArr f

    i img = let m = generate $ \(x, y) -> fromIntegral (pixelAt img x y) in force m

    f :: Diff (Matrix 28 28) (Vector 784)
    f = Diff $ \m -> generate $ \w -> m ! (w `mod` 28, w `div` 28)