Add tf.conv3dTranspose (#1629)

this PR proposes adding a conv3dTranspose op (with tests included).

This PR is driven by my group's desire to use a pre-trained 3D U-Net with Tensorflow JS.

FEATURE

Author: Jakub Kaczmarzyk

.vscode/settings.json

@@ -18,6 +18,7 @@
   "[javascript]": {
     "editor.formatOnSave": true
   },
+  "editor.defaultFormatter": "xaver.clang-format",
   "editor.rulers": [80],
   "clang-format.style": "Google",
   "files.insertFinalNewline": true,

src/ops/conv.ts

@@ -937,6 +937,36 @@ function conv3dDerFilter_<T extends Tensor4D|Tensor5D>(
       backend => backend.conv3dDerFilter(x5D, dy5D, convInfo), {x5D, dy5D});
 }
 
+/**
+ * Computes the transposed 3D convolution of a volume, also known as a
+ * deconvolution.
+ *
+ * @param x The input image, of rank 5 or rank 4, of shape
+ *   `[batch, depth, height, width, inDepth]`. If rank 4, batch of 1 is assumed.
+ * @param filter The filter, rank 4, of shape
+ *     `[depth, filterHeight, filterWidth, outDepth, inDepth]`.
+ *     `inDepth` must match `inDepth` in `x`.
+ * @param outputShape Output shape, of rank 5 or rank 4:
+ *     `[batch, depth, height, width, outDepth]`. If rank 3, batch of 1 is
+ *    assumed.
+ * @param strides The strides of the original convolution:
+ *     `[strideDepth, strideHeight, strideWidth]`.
+ * @param pad  The type of padding algorithm used in the non-transpose version
+ *    of the op.
+ */
+/** @doc {heading: 'Operations', subheading: 'Convolution'} */
+function conv3dTranspose_<T extends Tensor4D|Tensor5D>(
+    x: T|TensorLike, filter: Tensor5D|TensorLike,
+    outputShape:
+        [number, number, number, number,
+         number]|[number, number, number, number],
+    strides: [number, number, number]|number, pad: 'valid'|'same'): T {
+  const $x = convertToTensor(x, 'x', 'conv3dTranspose');
+  const $filter = convertToTensor(filter, 'filter', 'conv3dTranspose');
+
+  return conv3dDerInput_(outputShape, $x, $filter, strides, pad);
+}
+
 export const conv1d = op({conv1d_});
 export const conv2d = op({conv2d_});
 export const conv3d = op({conv3d_});
@@ -945,3 +975,4 @@ export const conv2dDerInput = op({conv2dDerInput_});
 export const depthwiseConv2d = op({depthwiseConv2d_});
 export const separableConv2d = op({separableConv2d_});
 export const conv2dTranspose = op({conv2dTranspose_});
+export const conv3dTranspose = op({conv3dTranspose_});

src/ops/conv3d_transpose_test.ts

@@ -0,0 +1,85 @@
+/**
+ * @license
+ * 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 tf from '../index';
+import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
+import {expectArraysClose} from '../test_util';
+
+describeWithFlags('conv3dTranspose', ALL_ENVS, () => {
+  // Reference Python TensorFlow code
+  // ```python
+  // import numpy as np
+  // import tensorflow as tf
+  // tf.enable_eager_execution()
+  // x = np.array([2], dtype = np.float32).reshape(1, 1, 1, 1, 1)
+  // w = np.array([5, 4, 8, 7, 1, 2, 6, 3], dtype = np.float32).reshape(2, 2, 2,
+  //   1, 1)
+  // tf.nn.conv3d_transpose(x, w, output_shape=[1, 2, 2, 2, 1], padding='VALID')
+  // ```
+  it('input=2x2x2x1,d2=1,f=2,s=1,p=valid', async () => {
+    const origInputDepth = 1;
+    const origOutputDepth = 1;
+    const inputShape: [number, number, number, number] =
+        [1, 1, 1, origOutputDepth];
+    const fSize = 2;
+    const origPad = 'valid';
+    const origStride = 1;
+
+    const x = tf.tensor4d([2], inputShape);
+    const w = tf.tensor5d(
+        [5, 4, 8, 7, 1, 2, 6, 3],
+        [fSize, fSize, fSize, origInputDepth, origOutputDepth]);
+
+    const result = tf.conv3dTranspose(x, w, [2, 2, 2, 1], origStride, origPad);
+    const expected = [10, 8, 16, 14, 2, 4, 12, 6];
+
+    expect(result.shape).toEqual([2, 2, 2, 1]);
+    expectArraysClose(await result.data(), expected);
+  });
+
+  // Reference Python TensorFlow code
+  // ```python
+  // import numpy as np
+  // import tensorflow as tf
+  // tf.enable_eager_execution()
+  // x = np.array([2, 3], dtype = np.float32).reshape(2, 1, 1, 1, 1, 1)
+  // w = np.array([5, 4, 8, 7, 1, 2, 6, 3], dtype = np.float32).reshape(2,
+  //   2, 2, 1, 1)
+  // tf.nn.conv3d_transpose(x, w, output_shape=[2, 2, 2, 2, 1], padding='VALID')
+  // ```
+  it('input=2x2x2x1,d2=1,f=2,s=1,p=valid, batch=2', async () => {
+    const origInputDepth = 1;
+    const origOutputDepth = 1;
+    const inputShape: [number, number, number, number, number] =
+        [2, 1, 1, 1, origOutputDepth];
+    const fSize = 2;
+    const origPad = 'valid';
+    const origStride = 1;
+
+    const x = tf.tensor5d([2, 3], inputShape);
+    const w = tf.tensor5d(
+        [5, 4, 8, 7, 1, 2, 6, 3],
+        [fSize, fSize, fSize, origInputDepth, origOutputDepth]);
+
+    const result =
+        tf.conv3dTranspose(x, w, [2, 2, 2, 2, 1], origStride, origPad);
+    const expected = [10, 8, 16, 14, 2, 4, 12, 6, 15, 12, 24, 21, 3, 6, 18, 9];
+
+    expect(result.shape).toEqual([2, 2, 2, 2, 1]);
+    expectArraysClose(await result.data(), expected);
+  });
+});

src/tests.ts

@@ -55,6 +55,7 @@ import './ops/conv2d_separable_test';
 import './ops/conv2d_test';
 import './ops/conv2d_transpose_test';
 import './ops/conv3d_test';
+import './ops/conv3d_transpose_test';
 import './ops/conv_util_test';
 import './ops/diag_test';
 import './ops/dropout_test';