Higher level ops for building neural network layers with deeplearn.js

layers.ts

import {Array1D, NDArray, NDArrayMathGPU, Scalar, Graph, InCPUMemoryShuffledInputProviderBuilder, Session, FeedEntry, AdamOptimizer, Initializer, ZerosInitializer, VarianceScalingInitializer, CostReduction} from 'deeplearn';
import { Tensor } from 'deeplearn/dist/src/graph/graph';

const learning_rate = 0.001
const num_steps = 2000
const batch_size = 128

const num_classes = 10 // MNIST total classes (0-9 digits)
const num_input = 784 // MNIST data input (img shape: 28*28)


const math = new NDArrayMathGPU();
const g = new Graph();
var session = new Session(g, math);

const xhr = new XMLHttpRequest();
xhr.open('GET', 'data/mnist.json');
xhr.onload = async () => {
  const data = JSON.parse(xhr.responseText);
  const images = data.images.map((x) => { return Array1D.new(x).reshape([28, 28, 1]) });
  const labels = data.labels.map((x) => { return Array1D.new(x) });
  console.log(data.images[0].reduce((a, b) => a + b, 0))

  const inputProviderBuilder = new InCPUMemoryShuffledInputProviderBuilder([images, labels]);
  const [imageProvider, labelProvider] = inputProviderBuilder.getInputProviders();

  var x: Tensor = g.placeholder("x", [28, 28, 1])
  var y: Tensor = g.placeholder("y", [num_classes])

  const feedEntries: FeedEntry[] = [
    {tensor: x, data: imageProvider},
    {tensor: y, data: labelProvider}
  ];

  var network = conv2d(x, 32, 5, g, undefined, undefined, undefined, (layer, graph) => {return graph.relu(layer)});
  network = g.maxPool(network, 2, 1, 0);
  var network = conv2d(network, 64, 3, g, undefined, undefined, undefined, (layer, graph) => {return graph.relu(layer)});
  network = g.maxPool(network, 2, 1, 0);
  network = flatten(network);
  network = g.layers.dense("dense", network, num_classes);

  const optimizer = new AdamOptimizer(learning_rate, 0.9, 0.999)
  const costTensor = g.softmaxCrossEntropyCost(network, y);

  for (let i = 0; i < 10; i++) {
    await math.scope(async () => {
      const cost = session.train(
        costTensor, feedEntries, batch_size, optimizer, CostReduction.MEAN);
      console.log(`last average cost (): ${await cost.val()}`);
    });
  }
}
xhr.onerror = (err) => console.error(err);
xhr.send();

/**
   * Functional interface for the 2D convolution layer.
   * @param inputs Tensor input.
   * @param filters Integer, the dimensionality of the output space (i.e. the number of filters in the convolution).
   * @param kernel_size Number to specify the height and width of the 2D convolution window.
   * @param graph Graph opbject.
   * @param strides Number to specify the strides of convolution.
   * @param padding One of "valid" or "same" (case-insensitive).
   * @param data_format "channels_last" or "channel_first"
   * @param activation Optional. Activation function which is applied on the final layer of the function. Function should accept Tensor and graph as parameters
   * @param kernel_initializer An initializer object for the convolution kernel.
   * @param bias_initializer  An initializer object for bias.
   * @param name string which represents name of the layer.
   * @return Tensor object.
   */

function conv2d(
  inputs: Tensor, 
  filters: number, 
  kernel_size: number, 
  graph: Graph, 
  strides: number = 1, 
  padding = "valid", 
  data_format = "channels_last", 
  activation?, 
  kernel_initializer: Initializer = new VarianceScalingInitializer(), 
  bias_initializer: Initializer = new ZerosInitializer(),
  name: string = "") {

  const channel_axis = data_format == "channels_last" ? inputs.shape[2] : inputs.shape[0];
  const depthwise_kernel_shape = [kernel_size, kernel_size, channel_axis, filters];

  var weights = graph.variable(name + "w", kernel_initializer.initialize(depthwise_kernel_shape, kernel_size * kernel_size * channel_axis * filters, filters))

  var bias = graph.variable(name + "b", bias_initializer.initialize([filters], kernel_size, filters))

  const layer = graph.conv2d(inputs, weights, bias, kernel_size, filters, strides, padding == "valid" || padding == "VALID" ? 0 : undefined);

  return activation == undefined ? layer : activation(layer, graph);
}

function flatten(network: Tensor) {
  return g.reshape(network, (() => {
    let i = 1;
    network.shape.forEach((val) => { i *= val });
    return [i];
  })());
}