cnn.tw

# cnn.tw: a convolutional neural network, trained end-to-end with `grad`.
#
# The task: tell a vertical bar from a horizontal bar in a noisy 6x6 image. The
# model is a real conv net (conv -> relu -> max-pool -> dense -> sigmoid) and
# `grad` differentiates the whole thing at once, convolution included. Same
# autodiff, static shapes, and single binary as the rest of Twill.

import "std/nn" as nn

seed(0)

# Build images procedurally: a vertical bar is a one-hot column broadcast down;
# a horizontal bar is a one-hot row broadcast across. Then add a little noise.
fn vbar(c) = ones(6, 1) @ reshape(nn.onehot(c, 6), 1, 6)
fn hbar(r) = reshape(nn.onehot(r, 6), 6, 1) @ ones(1, 6)
fn noisy(img) = img + 0.15 * randn(6, 6)

let images = [noisy(vbar(1)), noisy(vbar(3)), noisy(vbar(4)),
              noisy(hbar(1)), noisy(hbar(2)), noisy(hbar(4))]
let labels = [0.0, 0.0, 0.0, 1.0, 1.0, 1.0]
let N = 6

# Model: one conv layer (3 filters, 3x3, kernel [3,1,3,3] + bias [3,1,1]) then a
# dense head. After conv+pool each image is 3 channels of 2x2 = 12 features.
let model = {
  c1: nn.conv_init(3, 1, 3, 3),
  w: nn.he_init(1, 12),
  b: zeros(1)
}

# conv (6x6 -> 4x4) -> relu -> maxpool (-> 2x2) -> flatten -> dense -> sigmoid.
fn forward(m, img) {
  let x = reshape(img, 1, 6, 6)
  let feat = maxpool2d(relu(nn.conv(m.c1, x)), 2)   # [3, 2, 2]
  sigmoid(nn.dense(m.w, m.b, reshape(feat, 12)))
}

fn loss(m) {
  let total = 0.0
  for i in range(N) {
    total = total + nn.bce(forward(m, images[i]), labels[i])
  }
  total / N
}

# Train with Adam. The optimizer walks the model's leaves (including the conv
# kernel nested in a list), so there's no per-parameter update to hand-write.
let mo = nn.zeros_like(model)
let vo = nn.zeros_like(model)
for step in range(300) {
  let g = grad(loss)(model)
  let out = nn.adam_step(model, g, mo, vo, step + 1, 0.01, 0.9, 0.999, 0.00000001)
  model = out[0]
  mo = out[1]
  vo = out[2]
  if step % 60 == 0 {
    print("step", step, "loss", loss(model))
  }
}

let correct = 0.0
for i in range(N) {
  let p = item(forward(model, images[i]))
  if (p > 0.5) == (labels[i] > 0.5) {
    correct = correct + 1.0
  }
}
print("final loss:", loss(model))
print("accuracy:", correct / N)