import plot

def relu (input : Float) : Float = select (input > 0.0) input 0.0

instance [Add a, Add b] Add (a & b) add = \(a, b) (c, d). ( (a + c), (b + d)) sub = \(a, b) (c, d). ( (a - c), (b - d)) zero = (zero, zero)

instance [VSpace a, VSpace b] VSpace (a & b) scaleVec = \ s (a, b) . (scaleVec s a, scaleVec s b)

data Layer inp:Type out:Type params:Type = AsLayer {forward:(params -> inp -> out) & init:(Key -> params)}

def forward (l:Layer i o p) (p : p) (x : i): o = (AsLayer l' ) = l (getAt #forward l') p x

def init (l:Layer i o p) (k:Key) : p = (AsLayer l') = l (getAt #init l') k

def DenseParams (a:Type) (b:Type) : Type = ((a=>b=>Float) & (b=>Float))

def dense (a:Type) (b:Type) : Layer (a=>Float) (b=>Float) (DenseParams a b) = AsLayer { forward = (\ ((weight, bias)) x . for j. (bias.j + sum for i. weight.i.j * x.i)), init = arb }

def CNNParams (inc:Type) (outc:Type) (kw:Int) (kh:Int) : Type = ((outc=>inc=>Fin kh=>Fin kw=>Float) & (outc=>Float))

def conv2d (x:inc=>(Fin h)=>(Fin w)=>Float) (kernel:outc=>inc=>(Fin kh)=>(Fin kw)=>Float) : outc=>(Fin h)=>(Fin w)=>Float = for o i j. (i', j') = (ordinal i, ordinal j) case (i' + kh) <= h && (j' + kw) <= w of True -> sum for (ki, kj, inp). (di, dj) = (fromOrdinal (Fin h) (i' + (ordinal ki)), fromOrdinal (Fin w) (j' + (ordinal kj))) x.inp.di.dj * kernel.o.inp.ki.kj False -> zero

def cnn (h:Int) ?-> (w:Int) ?-> (inc:Type) (outc:Type) (kw:Int) (kh:Int) : Layer (inc=>(Fin h)=>(Fin w)=>Float) (outc=>(Fin h)=>(Fin w)=>Float) (CNNParams inc outc kw kh) = AsLayer { forward = (\ (weight, bias) x. for o i j . (conv2d x weight).o.i.j + bias.o), init = arb }

def split (x: m=>v) : n=>o=>v = for i j. x.((ordinal (i,j))@m)

def imtile (x: a=>b=>v) : n=>o=>p=>q=>v = for kw kh w h. (split (split x).w.kw).h.kh

def meanpool (kh: Type) (kw: Type) (x : m=>n=> Float) : ( h=>w=> Float) = out : (kh => kw => h => w => Float) = imtile x mean for (i,j). out.i.j

[k1, k2] = splitKey $ newKey 1

x1 : Fin 100 => Float = arb k1

x2 : Fin 100 => Float = arb k2

y = for i. case ((x1.i > 0.0) && (x2.i > 0.0)) || ((x1.i < 0.0) && (x2.i < 0.0)) of True -> 1 False -> 0

xs = for i. [x1.i, x2.i]

:html showPlot $ xycPlot x1 x2 $ for i. IToF y.i

simple = \h1. ndense1 = dense (Fin 2) h1 ndense2 = dense h1 (Fin 2) AsLayer { forward = (\ (dense1, dense2) x. x1' = forward ndense1 dense1 x x1 = for i. relu x1'.i logsoftmax $ forward ndense2 dense2 x1), init = (\key. [k1, k2] = splitKey key (init ndense1 k1, init ndense2 k2)) }

:t simple

def trainClass [VSpace p] (model: Layer a (b=>Float) p) (x: batch=>a) (y: batch=>b) (epochs : Type) (minibatch : Type) (minibatches : Type) : (epochs => p & epochs => Float ) = xs : minibatches => minibatch => a = split x ys : minibatches => minibatch => b = split y unzip $ withState (init model $ newKey 0) $ \params . for _ : epochs. loss = sum $ for b : minibatches. (loss, gradfn) = vjp (\ params. -sum for j. result = forward model params xs.b.j result.(ys.b.j)) (get params) gparams = gradfn 1.0 params := (get params) - scaleVec (0.05 / (IToF 100)) gparams loss (get params, loss)

simple_model = simple (Fin 10)

(all_params,losses) = trainClass simple_model xs (for i. (y.i @ (Fin 2))) (Fin 500) (Fin 100) (Fin 1)

span = linspace (Fin 10) (-1.0) (1.0)

tests = for h : (Fin 50). for i . for j. r = forward simple_model all_params.((ordinal h * 10)@_) [span.i, span.j] [exp r.(1@_), exp r.(0@_), 0.0]

:html imseqshow tests

H = 28

W = 28

Image = Fin 1 => Fin H => Fin W => Float

Class = Fin 10

lenet = \h1 h2 h3 . ncnn1 = cnn (Fin 1) h1 3 3 ncnn2 = cnn h1 h2 3 3 Pooled = (h2 & Fin 7 & Fin 7) ndense1 = dense Pooled h3 ndense2 = dense h3 Class AsLayer { forward = (\ (cnn1, cnn2, dense1, dense2) inp. x:Image = inp x1' = forward ncnn1 cnn1 x x1 = for i j k. relu x1'.i.j.k x2' = forward ncnn2 cnn2 x1 x2 = for i j k. relu x2'.i.j.k x3 : (h2 => Fin 7 => Fin 7 => Float) = for c. meanpool (Fin 4) (Fin 4) x2.c x4' = forward ndense1 dense1 for (i,j,k). x3.i.j.k x4 = for i. relu x4'.i logsoftmax $ forward ndense2 dense2 x4), init = (\key. [k1, k2, k3, k4] = splitKey key (init ncnn1 k1, init ncnn2 k2, init ndense1 k3, init ndense2 k4)) }

:t lenet

Batch = Fin 5000

Full = Fin ((size Batch) * H * W)

def pixel (x:Char) : Float32 = r = W8ToI x IToF case r < 0 of True -> (abs r) + 128 False -> r

def getIm : Batch => Image = (AsList _ im) = unsafeIO do readFile "examples/mnist.bin" raw = unsafeCastTable Full im for b: Batch c: (Fin 1) i:(Fin W) j:(Fin H). pixel raw.((ordinal (b, i, j)) @ Full)

def getLabel : Batch => Class = (AsList _ im2) = unsafeIO do readFile "examples/labels.bin" r = unsafeCastTable Batch im2 for i. (W8ToI r.i @ Class)

ims = getIm

labels = getLabel

small_ims = for i: (Fin 10). ims.((ordinal i)@_)

small_labels = for i: (Fin 10). labels.((ordinal i)@_)

:p small_labels

Epochs = (Fin 5)

Minibatches = (Fin 1)

Minibatch = (Fin 10)

:t ims.(2@_)

model = lenet (Fin 1) (Fin 1) (Fin 20)

init_param = (init model $ newKey 0)

:p forward model init_param (ims.(2@Batch))

:t (grad ((\x param. sum (forward model param x)) (ims.(2@_)))) init_param

(all_params', losses') = trainClass model small_ims small_labels Epochs Minibatch Minibatches

:p losses'