rework_mnist_arraymancer.nim

import ../src/arraymancer, random

randomize(42)

let
  ctx = newContext Tensor[float32] # Autograd/neural network graph
  n = 32                           # Batch size

let
  x_train = read_mnist_images("build/train-images.idx3-ubyte").astype(float32) / 255'f32
  X_train = ctx.variable x_train.unsqueeze(1)
  y_train = read_mnist_labels("build/train-labels.idx1-ubyte").astype(int)
  x_test = read_mnist_images("build/t10k-images.idx3-ubyte").astype(float32) / 255'f32
  X_test = ctx.variable x_test.unsqueeze(1)
  y_test = read_mnist_labels("build/t10k-labels.idx1-ubyte").astype(int)

type
  TrainableLayer[T] = object {.inheritable.}
    weight: Variable[T]
    bias: Variable[T]

  Conv2DLayer{.final.}[T] = object of TrainableLayer[T]

  LinearLayer{.final.}[T] = object of TrainableLayer[T]

type ExampleModel[TT] = object
  ## Holds the Trainable layers
  cv1: Conv2DLayer[TT]
  cv2: Conv2DLayer[TT]
  fc3: LinearLayer[TT]
  classifier: LinearLayer[TT]

proc init[TT](model_type: typedesc[ExampleModel[TT]]): ExampleModel[TT] =

  result.cv1.weight = ctx.variable(randomTensor(20, 1, 5, 5, 1'f32) .- 0.5'f32, requires_grad = true)
  result.cv1.bias   = ctx.variable(randomTensor(20, 1, 1, 1'f32) .- 0.5'f32, requires_grad = true)

  result.cv2.weight = ctx.variable(randomTensor(50, 20, 5, 5, 1'f32) .- 0.5'f32, requires_grad = true)
  result.cv2.bias   = ctx.variable(randomTensor(50, 1, 1, 1'f32) .- 0.5'f32, requires_grad = true)

  result.fc3.weight = ctx.variable(randomTensor(500, 800, 1'f32) .- 0.5'f32, requires_grad = true)
  result.classifier.weight = ctx.variable(randomTensor(10, 500, 1'f32) .- 0.5'f32, requires_grad = true)

proc forward[TT](self: ExampleModel[TT], x: Variable[TT]): Variable[TT] =

  template cv1(x: Variable): Variable = x.conv2d(self.cv1.weight, self.cv1.bias)
  template mp1(x: Variable): Variable = x.maxpool2D((2,2), (0,0), (2,2))
  template cv2(x: Variable): Variable = x.conv2d(self.cv2.weight, self.cv2.bias)
  template mp2(x: Variable): Variable = x.maxpool2D((2,2), (0,0), (2,2))
  template fc3(x: Variable): Variable = x.linear(self.fc3.weight)
  template classifier(x: Variable): Variable = x.linear(self.classifier.weight)

  result = x.cv1.relu.mp1.cv2.mp2.flatten.fc3.relu.classifier

let model = init(ExampleModel[Tensor[float32]])

let optim = newSGD[float32](
  model.cv1.weight, model.cv1.bias, model.cv2.weight, model.cv2.bias, model.fc3.weight, model.classifier.weight, 0.01f
)

# Learning loop
for epoch in 0 ..< 5:
  for batch_id in 0 ..< X_train.value.shape[0] div n:
    let offset = batch_id * n
    let x = X_train[offset ..< offset + n, _]
    let target = y_train[offset ..< offset + n]

    let clf = model.forward(x)
    let loss = clf.sparse_softmax_cross_entropy(target)

    if batch_id mod 200 == 0:
      echo "Epoch is: " & $epoch
      echo "Batch id: " & $batch_id
      echo "Loss is:  " & $loss.value.data[0]

    loss.backprop()
    optim.update()

  ctx.no_grad_mode:
    echo "\nEpoch #" & $epoch & " done. Testing accuracy"

    var score = 0.0
    var loss = 0.0
    for i in 0 ..< 10:
      let y_pred = model.forward(X_test[i*1000 ..< (i+1)*1000, _]).value.softmax.argmax(axis = 1).indices.squeeze
      score += accuracy_score(y_test[i*1000 ..< (i+1)*1000], y_pred)

      loss += model.forward(X_test[i*1000 ..< (i+1)*1000, _]).sparse_softmax_cross_entropy(y_test[i*1000 ..< (i+1)*1000]).value.data[0]
    score /= 10
    loss /= 10
    echo "Accuracy: " & $(score * 100) & "%"
    echo "Loss:     " & $loss
    echo "\n"