Change conv2d to match TF API (#67)

* update the readme and npm package

* update readme

* change conv2d api to take padding same|valid

* Merge branch 'master' into conv

* move logic to math

* some progress

* Merge remote-tracking branch 'origin/master' into conv

* Migrate all conv-related ops

* rename conv2dTranspose to conv2dDerInput

* switch shader indexing from float to int

* revert graph_runner_test

* self review

* Merge remote-tracking branch 'origin/master' into conv

* merge with the branch int_indexing

* fix typos in shaders

* merge with master

* update pool ops

* remove commented out code

* simplify api

* added unit tests for conv/pool

* add doc

* Merge master into conv

Author: dsmilkov

demos/benchmarks/conv_gpu_benchmark.ts

@@ -29,22 +29,22 @@ export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
   const texManager = new TextureManager(gpgpu);
   initializeGPU(gpgpu, texManager);
 
-  const inputDepth = 1;
-  const inputShape: [number, number, number] = [size, size, inputDepth];
-  const outputDepth = 1;
-  const fieldSize = 11;
+  const inDepth = 1;
+  const inShape: [number, number, number] = [size, size, inDepth];
+  const outDepth = 1;
+  const filterSize = 11;
   const stride = 1;
-  const zeroPad = conv_util.computeDefaultPad(inputShape, fieldSize, stride);
-
   const hasBias = true;
-  const program = new Conv2DProgram(
-      inputShape, fieldSize, outputDepth, stride, zeroPad, hasBias);
+  const convInfo = conv_util.computeConvInfo(
+      inShape, filterSize, filterSize, outDepth, stride, stride, 'same');
+  const program = new Conv2DProgram(convInfo, hasBias);
   const outputShape = program.outputShape as [number, number, number];
   const out = Array3D.zeros(outputShape);
-  const x = Array3D.randUniform(inputShape, -1, 1);
-  const wShape = conv_util.computeWeightsShape4D(1, outputDepth, fieldSize);
+  const x = Array3D.randUniform(inShape, -1, 1);
+  const wShape =
+      conv_util.computeWeightsShape4D(1, outDepth, filterSize, filterSize);
   const W = Array4D.randUniform(wShape, -1, 1);
-  const b = Array1D.randUniform([outputDepth], -1, 1);
+  const b = Array1D.randUniform([outDepth], -1, 1);
   const inputs = [x, W, b];
   const binary = gpgpu_math.compileProgram(gpgpu, program, inputs, out);
 

demos/benchmarks/conv_transpose_gpu_benchmark.ts

@@ -15,7 +15,7 @@ limitations under the License.
 
 import * as conv_util from '../../src/math/conv_util';
 import {Array3D, Array4D, initializeGPU} from '../../src/math/ndarray';
-import {Conv2DTransposeProgram} from '../../src/math/webgl/conv_backprop_gpu';
+import {Conv2DDerInputProgram} from '../../src/math/webgl/conv_backprop_gpu';
 import {GPGPUContext} from '../../src/math/webgl/gpgpu_context';
 import * as gpgpu_math from '../../src/math/webgl/gpgpu_math';
 import {TextureManager} from '../../src/math/webgl/texture_manager';
@@ -25,8 +25,8 @@ const OP_RUNS = 40;
 
 export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
   const origInputDepth = 1;
-  const origOutputDepth = 2;
-  const xShape: [number, number, number] = [size, size, 1];
+  const origOutputDepth = 1;
+  const xShape: [number, number, number] = [size, size, origOutputDepth];
   const fieldSize = 11;
   const origStride = 1;
   const origPad = 1;
@@ -36,14 +36,15 @@ export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
   initializeGPU(gpgpu, texManager);
   gpgpu.enableAutomaticDebugValidation(true);
 
-  const hasBias = false;
-  const program = new Conv2DTransposeProgram(
-      xShape, fieldSize, origInputDepth, origStride, origPad, hasBias);
+  const convInfo = conv_util.computeConvInfo(
+      xShape, fieldSize, fieldSize, origOutputDepth, origStride, origStride,
+      origPad);
+  const program = new Conv2DDerInputProgram(convInfo);
   const outputShape = program.outputShape as [number, number, number];
   const out = Array3D.zeros(outputShape);
   const x = Array3D.randUniform(xShape, -1, 1);
   const wShape = conv_util.computeWeightsShape4D(
-      origInputDepth, origOutputDepth, fieldSize);
+      origInputDepth, origOutputDepth, fieldSize, fieldSize);
   const W = Array4D.randUniform(wShape, -1, 1);
   const inputs = [x, W];
   const binary = gpgpu_math.compileProgram(gpgpu, program, inputs, out);

demos/benchmarks/max_pool_backprop_gpu_benchmark.ts

@@ -29,12 +29,14 @@ export const BENCHMARK_TEST: BenchmarkTest = (size: number) => {
   const texManager = new TextureManager(gpgpu);
   initializeGPU(gpgpu, texManager);
 
-  const outputDepth = 1;
-  const dyShape: [number, number, number] = [size, size, outputDepth];
+  const depth = 1;
+  const dyShape: [number, number, number] = [size, size, depth];
+  const xShape: [number, number, number] = [size, size, depth];
   const fSize = 11;
   const stride = 1;
-  const zeroPad = conv_util.computeDefaultPad(dyShape, fSize, stride);
-  const program = new MaxPool2DBackpropProgram(dyShape, fSize, stride, zeroPad);
+  const convInfo = conv_util.computeConvInfo(
+      xShape, fSize, fSize, depth, stride, stride, 'same');
+  const program = new MaxPool2DBackpropProgram(convInfo);
   const res = NDArray.zeros(program.outputShape);
   const dy = Array3D.randUniform(dyShape, -1, 1);
   const positionsData = new Float32Array(dy.size);

demos/benchmarks/max_pool_gpu_benchmark.ts

@@ -43,10 +43,9 @@ function testMaxPool(size: number, positions: boolean): number {
   const xShape: [number, number, number] = [size, size, outputDepth];
   const fieldSize = 11;
   const stride = 1;
-  const zeroPad = conv_util.computeDefaultPad(xShape, fieldSize, stride);
-
-  const program =
-      new Pool2DProgram(xShape, fieldSize, stride, zeroPad, 'max', positions);
+  const convInfo = conv_util.computeConvInfo(
+      xShape, fieldSize, fieldSize, outputDepth, stride, stride, 'same');
+  const program = new Pool2DProgram(convInfo, 'max', positions);
   const res = NDArray.zeros(program.outputShape);
   const x = Array3D.randUniform(xShape, -1, 1);
   const binary = gpgpu_math.compileProgram(gpgpu, program, [x], res);

demos/model-builder/layer_builder.ts

@@ -207,7 +207,7 @@ export class Convolution2DLayerBuilder implements LayerBuilder {
       {
         label: 'Output depth',
         initialValue: (inputShape: number[]) =>
-            this.outputDepth != null ? this.outputDepth : 1,
+                          this.outputDepth != null ? this.outputDepth : 1,
         type: 'number',
         min: 1,
         max: 1000,
@@ -319,7 +319,7 @@ export class ReshapeLayerBuilder implements LayerBuilder {
       initialValue: (inputShape: number[]) => inputShape.join(', '),
       type: 'text' as 'text',
       setValue: (value: string) => this.outputShape =
-          value.split(',').map((value) => +value),
+                    value.split(',').map((value) => +value),
       getValue: () => this.outputShape.join(', ')
     }];
   }

src/graph.ts

@@ -694,10 +694,9 @@ export class MaxPoolNode extends Node {
       graph: Graph, private x: Tensor, public fieldSize: number,
       public stride = 1, public zeroPad?: number) {
     super(
-        graph, 'Max pool', {x},
-        new Tensor(conv_util.computeOutputShape3D(
-            x.shape as [number, number, number], fieldSize, x.shape[2], stride,
-            zeroPad)));
+        graph, 'Max pool', {x}, new Tensor(conv_util.computeOutputShape3D(
+                                    x.shape as [number, number, number],
+                                    fieldSize, x.shape[2], stride, zeroPad)));
   }
   validate() {
     util.assert(
@@ -875,4 +874,4 @@ export class ArgMaxEqualsNode extends Node {
  * @hidden
  */
 export type ArrayData =
-    NDArray|number|number[]|number[][]|number[][][]|number[][][][];
+    NDArray | number | number[] | number[][] | number[][][] | number[][][][];

src/math/conv_util.ts

@@ -15,14 +15,88 @@ limitations under the License.
 
 import * as util from '../util';
 
+/**
+ * Information about the forward pass of a convolution/pooling operation.
+ * It includes input and output shape, strides, filter size and padding
+ * information.
+ */
+export type ConvInfo = {
+  inShape: [number, number, number],
+  outShape: [number, number, number],
+  strideHeight: number,
+  strideWidth: number,
+  filterHeight: number,
+  filterWidth: number,
+  padInfo: {top: number, left: number, right: number, bottom: number}
+};
+
+/**
+ * Computes the information about a forward pass of a convolution/pooling
+ * operation.
+ */
+export function computeConvInfo(
+    inShape: [number, number, number], filterHeight: number,
+    filterWidth: number, outDepth: number, strideHeight: number,
+    strideWidth: number, pad: 'same'|'valid'|number): ConvInfo {
+  if (typeof pad === 'number') {
+    const outShape = computeOutputShape3D(
+        inShape, filterHeight, outDepth, strideHeight, pad);
+    return {
+      inShape,
+      outShape,
+      padInfo: {top: pad, bottom: pad, left: pad, right: pad},
+      strideHeight,
+      strideWidth,
+      filterHeight,
+      filterWidth
+    };
+  }
+  const inHeight = inShape[0];
+  const inWidth = inShape[1];
+  let outShape: [number, number, number];
+  let padInfo: {left: number, top: number, bottom: number, right: number};
+  if (pad === 'same') {
+    const outHeight = Math.ceil(inHeight / strideHeight);
+    const outWidth = Math.ceil(inWidth / strideWidth);
+    outShape = [outHeight, outWidth, outDepth];
+    const padAlongHeight =
+        (outHeight - 1) * strideHeight + filterHeight - inHeight;
+    const padAlongWidth = (outWidth - 1) * strideWidth + filterWidth - inWidth;
+    const top = Math.floor(padAlongHeight / 2);
+    const bottom = padAlongHeight - top;
+    const left = Math.floor(padAlongWidth / 2);
+    const right = padAlongWidth - left;
+    padInfo = {top, bottom, left, right};
+  } else if (pad === 'valid') {
+    const outHeight = Math.ceil((inHeight - filterHeight + 1) / strideHeight);
+    const outWidth = Math.ceil((inWidth - filterWidth + 1) / strideWidth);
+    outShape = [outHeight, outWidth, outDepth];
+    padInfo = {top: 0, bottom: 0, left: 0, right: 0};
+  } else {
+    throw Error(`Unknown padding parameter: ${pad}`);
+  }
+  return {
+    inShape,
+    outShape,
+    padInfo,
+    strideHeight,
+    strideWidth,
+    filterHeight,
+    filterWidth
+  };
+}
+
+/**
+ * @deprecated Use `conv_util.computeConvInfo` instead.
+ */
 export function computeOutputShape3D(
-    inputShapeRowColDepth: [number, number, number], fieldSize: number,
-    depth: number, stride: number, zeroPad?: number): [number, number, number] {
+    inShape: [number, number, number], fieldSize: number, outDepth: number,
+    stride: number, zeroPad?: number): [number, number, number] {
   if (zeroPad == null) {
-    zeroPad = computeDefaultPad(inputShapeRowColDepth, fieldSize, stride);
+    zeroPad = computeDefaultPad(inShape, fieldSize, stride);
   }
-  const inputRows = inputShapeRowColDepth[0];
-  const inputCols = inputShapeRowColDepth[1];
+  const inputRows = inShape[0];
+  const inputCols = inShape[1];
   const outputRows = (inputRows - fieldSize + 2 * zeroPad) / stride + 1;
   util.assert(
       util.isInt(outputRows),
@@ -35,7 +109,7 @@ export function computeOutputShape3D(
       `The output # of columns (${outputCols}) must be an integer. Change ` +
           `the stride and/or zero pad parameters`);
 
-  return [outputRows, outputCols, depth];
+  return [outputRows, outputCols, outDepth];
 }
 
 export function computeDefaultPad(
@@ -50,9 +124,9 @@ export function computeTexShapeFrom3D(
 }
 
 export function computeWeightsShape4D(
-    inputDepth: number, outputDepth: number,
-    fSize: number): [number, number, number, number] {
-  return [fSize, fSize, inputDepth, outputDepth];
+    inputDepth: number, outputDepth: number, filterHeight: number,
+    filterWidth: number): [number, number, number, number] {
+  return [filterHeight, filterWidth, inputDepth, outputDepth];
 }
 
 export function computeDilatedRC(

src/math/conv_util_test.ts

@@ -0,0 +1,77 @@
+/* Copyright 2017 Google Inc. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+import * as conv_util from './conv_util';
+
+describe('conv_util computeConvInfo', () => {
+  it('1x1 conv over 1x1 array with same pad', () => {
+    const inShape: [number, number, number] = [1, 1, 1];
+    const convInfo = conv_util.computeConvInfo(inShape, 1, 1, 1, 1, 1, 'same');
+    expect(convInfo.outShape).toEqual([1, 1, 1]);
+  });
+
+  it('2x2 conv over 3x3 array with same pad', () => {
+    const inShape: [number, number, number] = [3, 3, 1];
+    const convInfo = conv_util.computeConvInfo(inShape, 2, 2, 1, 1, 1, 'same');
+    expect(convInfo.outShape).toEqual([3, 3, 1]);
+    // Should produce non-even padding with extra pixel at the right/bottom.
+    expect(convInfo.padInfo.left).toBe(0);
+    expect(convInfo.padInfo.right).toBe(1);
+    expect(convInfo.padInfo.top).toBe(0);
+    expect(convInfo.padInfo.bottom).toBe(1);
+  });
+
+  it('2x2 conv over 3x3 array with same pad', () => {
+    const inShape: [number, number, number] = [3, 3, 1];
+    const convInfo = conv_util.computeConvInfo(inShape, 2, 2, 1, 1, 1, 'same');
+    expect(convInfo.outShape).toEqual([3, 3, 1]);
+  });
+
+  it('2x2 conv over 3x3 array with valid pad', () => {
+    const inShape: [number, number, number] = [3, 3, 1];
+    const convInfo = conv_util.computeConvInfo(inShape, 2, 2, 1, 1, 1, 'valid');
+    expect(convInfo.outShape).toEqual([2, 2, 1]);
+  });
+
+  it('2x2 conv over 3x3 array with valid pad with stride 2', () => {
+    const inShape: [number, number, number] = [3, 3, 1];
+    const convInfo = conv_util.computeConvInfo(inShape, 2, 2, 1, 2, 2, 'valid');
+    expect(convInfo.outShape).toEqual([1, 1, 1]);
+  });
+
+  it('2x2 conv over 3x3 array with valid pad with stride 2', () => {
+    const inShape: [number, number, number] = [3, 3, 1];
+    const convInfo = conv_util.computeConvInfo(inShape, 2, 2, 1, 2, 2, 'valid');
+    expect(convInfo.outShape).toEqual([1, 1, 1]);
+  });
+
+  it('2x1 conv over 3x3 array with valid pad with stride 1', () => {
+    const inShape: [number, number, number] = [3, 3, 1];
+    const convInfo = conv_util.computeConvInfo(inShape, 2, 1, 1, 1, 1, 'valid');
+    expect(convInfo.outShape).toEqual([2, 3, 1]);
+  });
+
+  it('2x1 conv over 3x3 array with valid pad with strides h=2, w=1', () => {
+    const inShape: [number, number, number] = [3, 3, 1];
+    const convInfo = conv_util.computeConvInfo(inShape, 2, 1, 1, 2, 1, 'valid');
+    expect(convInfo.outShape).toEqual([1, 3, 1]);
+  });
+
+  it('1x2 conv over 3x3 array with valid pad with stride 1', () => {
+    const inShape: [number, number, number] = [3, 3, 1];
+    const convInfo = conv_util.computeConvInfo(inShape, 1, 2, 1, 1, 1, 'valid');
+    expect(convInfo.outShape).toEqual([3, 2, 1]);
+  });
+});