feat: pool & conv feedforward (#18)

* feat: neural network api

* feat: initial flatten, conv2d, pool2d layer structures

* feat: initial dropout layer

* chore(rust): rename backend => cpu

Co-authored-by: Dean Srebnik <[email protected]>

* feat(rust): conv2d init

* feat: WIP conv2d layer

* conv2d feedforward

* feat: max and avg pool api function

* fix: tensor types

* feat: conv2d wasm example

* feat: pool feedforward

Co-authored-by: Dean Srebnik <[email protected]>

---------

Co-authored-by: Dean Srebnik <[email protected]>
Co-authored-by: Dean Srebnik <[email protected]>

deps.ts

@@ -1,6 +1,3 @@
-export { WebGPUBackend } from "https://raw.githubusercontent.com/denosaurs/neo/e7f602cd218fbdb7418e0fca9ac4aba8d6c22d08/backend/webgpu/backend.ts";
-export { Core } from "https://raw.githubusercontent.com/denosaurs/neo/e7f602cd218fbdb7418e0fca9ac4aba8d6c22d08/backend/core/core.ts";
-export { WebGPUData } from "https://raw.githubusercontent.com/denosaurs/neo/e7f602cd218fbdb7418e0fca9ac4aba8d6c22d08/backend/webgpu/data.ts";
 export {
   CsvStream,
 } from "https://deno.land/[email protected]/csv/mod.ts";

examples/filters/conv.ts

@@ -0,0 +1,100 @@
+import {
+  Activation,
+  Conv2DLayer,
+  Cost,
+  CPU,
+  MaxPool2DLayer,
+  Rank,
+  Sequential,
+  setupBackend,
+  Tensor,
+  tensor4D,
+} from "../../mod.ts";
+import { decode } from "https://deno.land/x/[email protected]/mod.ts";
+import { createCanvas } from "https://deno.land/x/[email protected]/mod.ts";
+import { Layer } from "../../src/core/api/layer.ts";
+
+const canvas = createCanvas(600, 600);
+const ctx = canvas.getContext("2d");
+ctx.fillStyle = "white";
+ctx.fillRect(0, 0, 600, 600);
+
+const dim = 28;
+const kernel = [[[
+  [-1, 1, 0],
+  [-1, 1, 0],
+  [-1, 1, 0],
+]]];
+
+//Credit: Hashrock (https://github.com/hashrock)
+const img = decode(Deno.readFileSync("./examples/filters/deno.png")).image;
+const buffer = new Float32Array(dim * dim);
+for (let i = 0; i < dim * dim; i++) {
+  buffer[i] = img[i * 4];
+}
+
+await setupBackend(CPU);
+
+drawPixels(buffer, dim);
+
+const conv = await feedForward([
+  Conv2DLayer({
+    activation: Activation.Linear,
+    kernel: tensor4D(kernel),
+    kernelSize: [1, 1, 3, 3],
+    padding: 1,
+    strides: [1, 1],
+    unbiased: true,
+  }),
+]);
+
+drawPixels(conv.data, conv.shape[2], 280);
+
+const pool = await feedForward([
+  Conv2DLayer({
+    activation: Activation.Linear,
+    kernel: tensor4D(kernel),
+    kernelSize: [1, 1, 3, 3],
+    padding: 1,
+    strides: [1, 1],
+    unbiased: true,
+  }),
+  MaxPool2DLayer({ strides: [2, 2] }),
+]);
+
+drawPixels(pool.data, pool.shape[2], 0, 280, 2);
+
+async function feedForward(layers: Layer[]) {
+  const net = new Sequential({
+    size: [1, 1, dim, dim],
+    silent: true,
+    layers,
+    cost: Cost.MSE,
+  });
+
+  const data = new Tensor(buffer, [1, 1, dim, dim]);
+  return await net.predict(data) as Tensor<Rank.R4>;
+}
+
+function drawPixels(
+  buffer: Float32Array,
+  dim: number,
+  offsetX = 0,
+  offsetY = 0,
+  scale = 1,
+) {
+  for (let i = 0; i < dim; i++) {
+    for (let j = 0; j < dim; j++) {
+      const pixel = buffer[j * dim + i];
+      ctx.fillStyle = `rgb(${pixel}, ${pixel}, ${pixel})`;
+      ctx.fillRect(
+        i * 10 * scale + offsetX,
+        j * 10 * scale + offsetY,
+        10 * scale,
+        10 * scale,
+      );
+    }
+  }
+}
+
+await Deno.writeFile("./examples/filters/output.png", canvas.toBuffer());

examples/filters/conv_wasm.ts

@@ -0,0 +1,94 @@
+import {
+  Activation,
+  Conv2DLayer,
+  Cost,
+  MaxPool2DLayer,
+  Rank,
+  Sequential,
+  setupBackend,
+  Tensor,
+  tensor4D,
+  WASM,
+} from "../../mod.ts";
+import { decode } from "https://deno.land/x/[email protected]/mod.ts";
+import { createCanvas } from "https://deno.land/x/[email protected]/mod.ts";
+import { Layer } from "../../src/core/api/layer.ts";
+
+const canvas = createCanvas(600, 600);
+const ctx = canvas.getContext("2d");
+ctx.fillStyle = "white";
+ctx.fillRect(0, 0, 600, 600);
+
+const dim = 28;
+const kernel = [[[
+  [-1, 1, 0],
+  [-1, 1, 0],
+  [-1, 1, 0],
+]]];
+
+//Credit: Hashrock (https://github.com/hashrock)
+const img = decode(Deno.readFileSync("./examples/filters/deno.png")).image;
+const buffer = new Float32Array(dim * dim);
+for (let i = 0; i < dim * dim; i++) {
+  buffer[i] = img[i * 4];
+}
+
+await setupBackend(WASM);
+
+drawPixels(buffer, dim);
+
+const conv = await feedForward([
+  Conv2DLayer({
+    activation: Activation.Linear,
+    kernel: tensor4D(kernel),
+    kernelSize: [1, 1, 3, 3],
+    padding: 1,
+    strides: [1, 1],
+    unbiased: true,
+  }),
+]);
+
+drawPixels(conv.data, conv.shape[2]);
+
+const pool = await feedForward([
+  Conv2DLayer({
+    activation: Activation.Linear,
+    kernel: tensor4D(kernel),
+    kernelSize: [1, 1, 3, 3],
+    padding: 1,
+    strides: [1, 1],
+    unbiased: true,
+  }),
+  MaxPool2DLayer({ strides: [2, 2] }),
+]);
+
+drawPixels(pool.data, pool.shape[2]);
+
+async function feedForward(layers: Layer[]) {
+  const net = new Sequential({
+    size: [1, 1, dim, dim],
+    silent: true,
+    layers,
+    cost: Cost.MSE,
+  });
+
+  const data = new Tensor(buffer, [1, 1, dim, dim]);
+  return await net.predict(data) as Tensor<Rank.R4>;
+}
+
+function drawPixels(
+  buffer: Float32Array,
+  dim: number,
+  offsetX = 0,
+  offsetY = 0,
+) {
+  for (let i = 0; i < dim; i++) {
+    for (let j = 0; j < dim; j++) {
+      const pixel = buffer[j * dim + i];
+      ctx.fillStyle = `rgb(${pixel}, ${pixel}, ${pixel})`;
+      ctx.fillRect(i * 10 + offsetX, j * 10 + offsetY, 10, 10);
+    }
+  }
+}
+
+await Deno.writeFile("./examples/filters/output.png", canvas.toBuffer());

examples/mnist/.gitignore

@@ -0,0 +1 @@
+*.idx

examples/mnist/common.ts

@@ -0,0 +1,64 @@
+import { DataSet } from "../../src/core/types.ts";
+import { Tensor } from "../../mod.ts";
+
+export function assert(condition: boolean, message?: string) {
+  if (!condition) {
+    throw new Error(message);
+  }
+}
+
+export function loadDataset(
+  imagesFile: string,
+  labelsFile: string,
+  start: number,
+  end: number,
+) {
+  const images = Deno.readFileSync(new URL(imagesFile, import.meta.url));
+  const labels = Deno.readFileSync(new URL(labelsFile, import.meta.url));
+
+  const imageView = new DataView(images.buffer);
+  const labelView = new DataView(labels.buffer);
+
+  assert(imageView.getUint32(0) === 0x803, "Invalid image file");
+  assert(labelView.getUint32(0) === 0x801, "Invalid label file");
+
+  const count = imageView.getUint32(4);
+  assert(count === labelView.getUint32(4), "Image and label count mismatch");
+
+  const inputs: Float32Array[] = [];
+  let mean = 0;
+  let sd = 0;
+  for (let i = 0; i < count; i++) {
+    const input = new Float32Array(784);
+    for (let j = 0; j < 784; j++) {
+      input[j] = imageView.getUint8(16 + i * 784 + j);
+      mean += input[j];
+      sd += Math.pow(input[j], 2);
+    }
+    inputs.push(input);
+  }
+
+  mean /= count * 784;
+  sd /= count * 784;
+  sd -= Math.pow(mean, 2);
+  sd = Math.sqrt(sd);
+
+  const results: DataSet[] = [];
+
+  for (let i = start; i < end; i++) {
+    for (let j = 0; j < 784; j++) {
+      inputs[i][j] -= mean;
+      inputs[i][j] /= sd;
+    }
+
+    const outputs = new Float32Array(10);
+    outputs[labelView.getUint8(8 + i)] = 1;
+
+    results.push({
+      inputs: new Tensor(inputs[i], [28, 28, 1, 1]),
+      outputs: new Tensor(outputs, [10]),
+    });
+  }
+
+  return results;
+}

examples/mnist/download.ts

@@ -0,0 +1,29 @@
+async function download(url: string, to: string) {
+  console.log("Download", url);
+  const f = await Deno.open(new URL(to, import.meta.url), {
+    write: true,
+    create: true,
+  });
+  await fetch(url).then((response) => {
+    response.body!.pipeThrough(new DecompressionStream("gzip")).pipeTo(
+      f.writable,
+    );
+  });
+}
+
+await download(
+  "http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
+  "train-images.idx",
+);
+await download(
+  "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
+  "train-labels.idx",
+);
+await download(
+  "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
+  "test-images.idx",
+);
+await download(
+  "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz",
+  "test-labels.idx",
+);

examples/mnist/train.ts

@@ -0,0 +1,79 @@
+import {
+  Conv2DLayer,
+  Cost,
+  CPU,
+  DenseLayer,
+  FlattenLayer,
+  Init,
+  MaxPool2DLayer,
+  Rank,
+  ReluLayer,
+  Sequential,
+  setupBackend,
+  SoftmaxLayer,
+  Tensor,
+} from "../../mod.ts";
+import { loadDataset } from "./common.ts";
+
+await setupBackend(CPU);
+
+// training
+const network = new Sequential({
+  size: [1, 1, 28, 28],
+  silent: true,
+  layers: [
+    Conv2DLayer({ kernelSize: [5, 5, 1, 6], padding: 2 }),
+    ReluLayer(),
+    MaxPool2DLayer({ strides: [2, 2] }),
+    Conv2DLayer({ kernelSize: [5, 5, 6, 16] }),
+    ReluLayer(),
+    MaxPool2DLayer({ strides: [2, 2] }),
+    Conv2DLayer({ kernelSize: [5, 5, 16, 120] }),
+    ReluLayer(),
+    FlattenLayer({ size: [120] }),
+    DenseLayer({ size: [84], init: Init.Kaiming }),
+    ReluLayer(),
+    DenseLayer({ size: [10], init: Init.Kaiming }),
+    SoftmaxLayer(),
+  ],
+  cost: Cost.CrossEntropy,
+});
+
+console.log("Loading training dataset...");
+const trainSet = loadDataset("train-images.idx", "train-labels.idx", 0, 5000);
+
+const epochs = 1;
+console.log("Training (" + epochs + " epochs)...");
+const start = performance.now();
+network.train(trainSet, epochs, 0.01);
+console.log("Training complete!", performance.now() - start);
+
+
+
+// predicting
+
+const testSet = loadDataset("test-images.idx", "test-labels.idx", 0, 1000);
+testSet.map((_, i) => testSet[i].inputs.shape = [28, 28, 1]);
+
+function argmax(mat: Tensor<Rank>) {
+  let max = -Infinity;
+  let index = -1;
+  for (let i = 0; i < mat.data.length; i++) {
+    if (mat.data[i] > max) {
+      max = mat.data[i];
+      index = i;
+    }
+  }
+  return index;
+}
+
+const correct = testSet.filter(async (e) => {
+  const prediction = argmax(await network.predict(e.inputs as Tensor<Rank>));
+  const expected = argmax(e.outputs as Tensor<Rank>);
+  return prediction === expected;
+});
+
+console.log(`${correct.length} / ${testSet.length} correct`);
+console.log(
+  `accuracy: ${((correct.length / testSet.length) * 100).toFixed(2)}%`,
+);