2d atrous convolution and atrous depthwise convolution (#794)

* Implemented atrous 1d and 2d convolution.

* Updated backend gpu and cpu to do atrous convolution when there is a dilation rate.
* Refactored backend_cpu.conv2d to be consistent with backend_gpu.conv2d, and also to make the way it does dilated convolution consistent.
* Added tests for 1d and 2d convolution with dilation rates that show the effect on the filter when dilation rates are set.
* Updated computeDefaultPad to account for a dilation rate.

Implemented atrous convolution for depthwiseConv2D

* Modified cpu depthwiseConv2D logic to be similar to that on the GPU,
so that dilation can be easily applied.

Still need to add more tests for depthwise atrous convolution

Per feedback, changed order of parameters to match
Tensorflow api.  Added dataFormat parameter to conv1d, conv2d, and
depthwiseConv2d, but did not yet implement that parameter; it defaults
to a value and cannot be a different value until the functionality
is implemented

raising error when gradient is done with atrous convolution

* Made ordering of parameters consistent in convolution methods with new dilation
parameters

* added depthwise conv tests

Author: oveddan

src/kernels/backend_cpu.ts

@@ -863,6 +863,8 @@ export class MathBackendCPU implements KernelBackend {
   conv2d(x: Tensor4D, filter: Tensor4D, convInfo: Conv2DInfo): Tensor4D {
     const filterHeight = convInfo.filterHeight;
     const filterWidth = convInfo.filterWidth;
+    const dilationHeight = convInfo.dilationHeight;
+    const dilationWidth = convInfo.dilationWidth;
     const padLeft = convInfo.padInfo.left;
     const padTop = convInfo.padInfo.top;
     const y = ops.buffer<Rank.R4>(convInfo.outShape, x.dtype);
@@ -871,17 +873,24 @@ export class MathBackendCPU implements KernelBackend {
       for (let d2 = 0; d2 < convInfo.outChannels; ++d2) {
         for (let yR = 0; yR < convInfo.outHeight; ++yR) {
           const xRCorner = yR * convInfo.strideHeight - padLeft;
-          const xRMin = Math.max(0, xRCorner);
-          const xRMax = Math.min(convInfo.inHeight, filterHeight + xRCorner);
           for (let yC = 0; yC < convInfo.outWidth; ++yC) {
             const xCCorner = yC * convInfo.strideWidth - padTop;
-            const xCMin = Math.max(0, xCCorner);
-            const xCMax = Math.min(convInfo.inWidth, filterWidth + xCCorner);
+
             let dotProd = 0;
-            for (let xR = xRMin; xR < xRMax; ++xR) {
-              const wR = xR - xRCorner;
-              for (let xC = xCMin; xC < xCMax; ++xC) {
-                const wC = xC - xCCorner;
+            for (let wR = 0; wR < filterHeight; wR++) {
+              const xR = xRCorner + wR * dilationHeight;
+
+              if (xR < 0 || xR >= convInfo.inHeight) {
+                continue;
+              }
+
+              for (let wC = 0; wC < filterWidth; wC++) {
+                const xC = xCCorner + wC * dilationWidth;
+
+                if (xC < 0 || xC >= convInfo.inWidth) {
+                  continue;
+                }
+
                 for (let d1 = 0; d1 < convInfo.inChannels; ++d1) {
                   const pixel = x.get(b, xR, xC, d1);
                   const weight = filter.get(wR, wC, d1, d2);
@@ -989,6 +998,8 @@ export class MathBackendCPU implements KernelBackend {
       Tensor4D {
     const filterHeight = convInfo.filterHeight;
     const filterWidth = convInfo.filterWidth;
+    const dilationHeight = convInfo.dilationHeight;
+    const dilationWidth = convInfo.dilationWidth;
     const padLeft = convInfo.padInfo.left;
     const padTop = convInfo.padInfo.top;
     const chMul = convInfo.outChannels / convInfo.inChannels;
@@ -998,18 +1009,24 @@ export class MathBackendCPU implements KernelBackend {
       for (let d1 = 0; d1 < convInfo.inChannels; ++d1) {
         for (let yR = 0; yR < convInfo.outHeight; ++yR) {
           const xRCorner = yR * convInfo.strideHeight - padLeft;
-          const xRMin = Math.max(0, xRCorner);
-          const xRMax = Math.min(convInfo.inHeight, filterHeight + xRCorner);
           for (let yC = 0; yC < convInfo.outWidth; ++yC) {
             const xCCorner = yC * convInfo.strideWidth - padTop;
-            const xCMin = Math.max(0, xCCorner);
-            const xCMax = Math.min(convInfo.inWidth, filterWidth + xCCorner);
             for (let q = 0; q < chMul; ++q) {
               let dotProd = 0;
-              for (let xR = xRMin; xR < xRMax; ++xR) {
-                const wR = xR - xRCorner;
-                for (let xC = xCMin; xC < xCMax; ++xC) {
-                  const wC = xC - xCCorner;
+              for (let wR = 0; wR < filterHeight; ++wR) {
+                const xR = xRCorner + wR * dilationHeight;
+
+                if (xR < 0 || xR >= convInfo.inHeight) {
+                  continue;
+                }
+
+                for (let wC = 0; wC < filterWidth; ++wC) {
+                  const xC = xCCorner + wC * dilationWidth;
+
+                  if (xC < 0 || xC >= convInfo.inWidth) {
+                    continue;
+                  }
+
                   const pixel = x.get(b, xR, xC, d1);
                   const weight = filter.get(wR, wC, d1, q);
                   dotProd += pixel * weight;
@@ -1021,6 +1038,7 @@ export class MathBackendCPU implements KernelBackend {
         }
       }
     }
+
     return y.toTensor();
   }
 

src/kernels/webgl/conv_gpu.ts

@@ -29,6 +29,8 @@ export class Conv2DProgram implements GPGPUProgram {
     const padLeft = convInfo.padInfo.left;
     const strideHeight = convInfo.strideHeight;
     const strideWidth = convInfo.strideWidth;
+    const dilationHeight = convInfo.dilationHeight;
+    const dilationWidth = convInfo.dilationWidth;
     const filterHeight = convInfo.filterHeight;
     const filterWidth = convInfo.filterWidth;
 
@@ -52,14 +54,14 @@ export class Conv2DProgram implements GPGPUProgram {
         // ? = to be determined. : = across all values in that axis.
         float dotProd = 0.0;
         for (int wR = 0; wR < ${filterHeight}; wR++) {
-          int xR = xRCorner + wR;
+          int xR = xRCorner + wR * ${dilationHeight};
 
           if (xR < 0 || xR >= ${convInfo.inHeight}) {
             continue;
           }
 
           for (int wC = 0; wC < ${filterWidth}; wC++) {
-            int xC = xCCorner + wC;
+            int xC = xCCorner + wC * ${dilationWidth};
 
             if (xC < 0 || xC >= ${convInfo.inWidth}) {
               continue;

src/kernels/webgl/conv_gpu_depthwise.ts

@@ -32,6 +32,8 @@ export class DepthwiseConv2DProgram implements GPGPUProgram {
     const padLeft = convInfo.padInfo.left;
     const strideHeight = convInfo.strideHeight;
     const strideWidth = convInfo.strideWidth;
+    const dilationHeight = convInfo.dilationHeight;
+    const dilationWidth = convInfo.dilationWidth;
     const filterHeight = convInfo.filterHeight;
     const filterWidth = convInfo.filterWidth;
     const channelMul = convInfo.outChannels / convInfo.inChannels;
@@ -56,14 +58,14 @@ export class DepthwiseConv2DProgram implements GPGPUProgram {
         float dotProd = 0.0;
         // TODO(dsmilkov): Flatten the two for loops and vec4 the operations.
         for (int wR = 0; wR < ${filterHeight}; wR++) {
-          int xR = xRCorner + wR;
+          int xR = xRCorner + wR * ${dilationHeight};
 
           if (xR < 0 || xR >= ${xNumRows}) {
             continue;
           }
 
           for (int wC = 0; wC < ${filterWidth}; wC++) {
-            int xC = xCCorner + wC;
+            int xC = xCCorner + wC * ${dilationWidth};
 
             if (xC < 0 || xC >= ${xNumCols}) {
               continue;

src/ops/conv.ts

@@ -40,6 +40,12 @@ export class ConvOps {
    *   - For more info, see this guide:
    *     [https://www.tensorflow.org/api_guides/python/nn#Convolution](
    *          https://www.tensorflow.org/api_guides/python/nn#Convolution)
+   * @param dataFormat An optional string from "NWC", "NCW". Defaults to "NWC",
+   *     the data is stored in the order of [batch, in_width, in_channels]. Only
+   *     "NWC" is currently supported.
+   * @param dilation The dilation rate in which we sample input values in
+   *     atrous convolution. Defaults to `1`. If it is greater than 1, then
+   *     stride must be `1`.
    * @param dimRoundingMode The rounding mode used when computing output
    *     dimensions if pad is a number. If none is provided, it will not round
    *     and error if the output is of fractional size.
@@ -48,6 +54,7 @@ export class ConvOps {
   @operation
   static conv1d<T extends Tensor2D|Tensor3D>(
       input: T, filter: Tensor3D, stride: number, pad: 'valid'|'same'|number,
+      dataFormat: 'NWC'|'NCW' = 'NWC', dilation = 1,
       dimRoundingMode?: 'floor'|'round'|'ceil'): T {
     let input3D = input as Tensor3D;
     let reshapedTo3D = false;
@@ -74,15 +81,27 @@ export class ConvOps {
         input3D.shape[2] === filter.shape[1],
         `Error in conv1d: depth of input (${input3D.shape[2]}) must match  ` +
             `input depth for filter ${filter.shape[1]}.`);
+    util.assert(
+        eitherStridesOrDilationsAreOne(stride, dilation),
+        'Error in conv1D: Either stride or dilation must be 1.' +
+            `Got stride ${stride} and dilation '${dilation}'`);
+    util.assert(
+        dataFormat === 'NWC',
+        `Error in conv1d: got dataFormat of ${
+            dataFormat} but only NWC is currently supported.`);
 
     const filter4D =
         filter.as4D(1, filter.shape[0], filter.shape[1], filter.shape[2]);
     const input4D =
         input3D.as4D(input3D.shape[0], 1, input3D.shape[1], input3D.shape[2]);
     const strides: [number, number] = [1, stride];
+    const dilations: [number, number] = [1, dilation];
+
+    const conv2dDataFormat = 'NHWC';
 
-    const res =
-        ConvOps.conv2d(input4D, filter4D, strides, pad, dimRoundingMode);
+    const res = ConvOps.conv2d(
+        input4D, filter4D, strides, pad, conv2dDataFormat, dilations,
+        dimRoundingMode);
 
     if (reshapedTo3D) {
       return res.as2D(res.shape[2], res.shape[3]) as T;
@@ -108,6 +127,15 @@ export class ConvOps {
    *   - For more info, see this guide:
    *     [https://www.tensorflow.org/api_guides/python/nn#Convolution](
    *          https://www.tensorflow.org/api_guides/python/nn#Convolution)
+   * @param dataFormat: An optional string from: "NHWC", "NCHW". Defaults to
+   *     "NHWC". Specify the data format of the input and output data. With the
+   *     default format "NHWC", the data is stored in the order of: [batch,
+   *     height, width, channels]. Only "NHWC" is currently supported.
+   * @param dilations The dilation rates: `[dilationHeight, dilationWidth]`
+   *     in which we sample input values across the height and width dimensions
+   *     in atrous convolution. Defaults to `[1, 1]`. If `dilations` is a single
+   *     number, then `dilationHeight == dilationWidth`. If it is greater than
+   *     1, then all values of `strides` must be 1.
    * @param dimRoundingMode The rounding mode used when computing output
    *     dimensions if pad is a number. If none is provided, it will not round
    *     and error if the output is of fractional size.
@@ -116,7 +144,9 @@ export class ConvOps {
   @operation
   static conv2d<T extends Tensor3D|Tensor4D>(
       x: T, filter: Tensor4D, strides: [number, number]|number,
-      pad: 'valid'|'same'|number, dimRoundingMode?: 'floor'|'round'|'ceil'): T {
+      pad: 'valid'|'same'|number, dataFormat: 'NHWC'|'NCHW' = 'NHWC',
+      dilations: [number, number]|number = [1, 1],
+      dimRoundingMode?: 'floor'|'round'|'ceil'): T {
     let x4D = x as Tensor4D;
     let reshapedTo4D = false;
 
@@ -142,13 +172,24 @@ export class ConvOps {
         x4D.shape[3] === filter.shape[2],
         `Error in conv2d: depth of input (${x4D.shape[3]}) must match  ` +
             `input depth for filter ${filter.shape[2]}.`);
-
-    const dilations = 1;
+    util.assert(
+        eitherStridesOrDilationsAreOne(strides, dilations),
+        'Error in conv2D: Either strides or dilations must be 1.' +
+            `Got strides ${strides} and dilations '${dilations}'`);
+    util.assert(
+        dataFormat === 'NHWC',
+        `Error in conv2d: got dataFormat of ${
+            dataFormat} but only NHWC is currently supported.`);
 
     const convInfo = conv_util.computeConv2DInfo(
         x4D.shape, filter.shape, strides, dilations, pad, dimRoundingMode);
 
     const grad = (dy: Tensor4D) => {
+      util.assert(
+          tupleValuesAreOne(dilations),
+          'Error in gradient of conv2D: dilation rates greater than 1 are not' +
+              `yet supported in gradients. Got dilations '${dilations}'`);
+
       return {
         x: () => ConvOps.conv2dDerInput(x4D.shape, dy, filter, strides, pad),
         filter: () =>
@@ -375,9 +416,13 @@ export class ConvOps {
    *          https://www.tensorflow.org/api_guides/python/nn#Convolution)
    * @param dilations The dilation rates: `[dilationHeight, dilationWidth]`
    *     in which we sample input values across the height and width dimensions
-   *     in atrous convolution. Defaults to `[1, 1]`. If `dilations` is a single
+   *     in atrous convolution. Defaults to `[1, 1]`. If `rate` is a single
    *     number, then `dilationHeight == dilationWidth`. If it is greater than
    *     1, then all values of `strides` must be 1.
+   * @param dataFormat: An optional string from: "NHWC", "NCHW". Defaults to
+   *     "NHWC". Specify the data format of the input and output data. With the
+   *     default format "NHWC", the data is stored in the order of: [batch,
+   *     height, width, channels]. Only "NHWC" is currently supported.
    * @param dimRoundingMode The rounding mode used when computing output
    *     dimensions if pad is a number. If none is provided, it will not round
    *     and error if the output is of fractional size.
@@ -386,7 +431,8 @@ export class ConvOps {
   @operation
   static depthwiseConv2d<T extends Tensor3D|Tensor4D>(
       input: T, filter: Tensor4D, strides: [number, number]|number,
-      pad: 'valid'|'same'|number, dilations: [number, number]|number = [1, 1],
+      pad: 'valid'|'same'|number, dataFormat: 'NHWC'|'NCHW' = 'NHWC',
+      dilations: [number, number]|number = [1, 1],
       dimRoundingMode?: 'floor'|'round'|'ceil'): T {
     let input4D = input as Tensor4D;
     let reshapedTo4D = false;
@@ -410,11 +456,11 @@ export class ConvOps {
     if (dilations == null) {
       dilations = [1, 1];
     }
-    const [dilationHeight, dilationWidth] = parseTupleParam(dilations);
     util.assert(
-        dilationHeight === 1 && dilationWidth === 1,
-        'Error in depthwiseConv2D: dilation rates greater than 1 are not yet ' +
-            `supported. Got dilations '${dilations}'`);
+        eitherStridesOrDilationsAreOne(strides, dilations),
+        'Error in depthwiseConv2d: Either strides or dilations must be 1.' +
+            `Got strides ${strides} and dilations '${dilations}'`);
+
     if (dimRoundingMode != null) {
       util.assert(
           util.isInt(pad as number),
@@ -438,3 +484,14 @@ export class ConvOps {
 function parseTupleParam(param: number|[number, number]): [number, number] {
   return typeof param === 'number' ? [param, param] : param;
 }
+
+function tupleValuesAreOne(param: number|[number, number]): boolean {
+  const [dimA, dimB] = parseTupleParam(param);
+  return dimA === 1 && dimB === 1;
+}
+
+function eitherStridesOrDilationsAreOne(
+    strides: number|[number, number],
+    dilations: number|[number, number]): boolean {
+  return tupleValuesAreOne(strides) || tupleValuesAreOne(dilations);
+}