Add transposedConv2D (#596) (#672)

Rename `conv2DBackpropInput `to `transposedConv2D` and make it public.
`transposedConv2D` is consistent with TensorFlow and PyTorch.

Author: t-ae

Sources/TensorFlow/Layers/Convolutional.swift

@@ -437,7 +437,7 @@ public struct TransposedConv1D<Scalar: TensorFlowFloatingPoint>: Layer {
             stride + (filter.shape[0] * paddingIndex)
         let c = filter.shape[2]
         let newShape = Tensor<Int32>([Int32(batchSize), 1, Int32(w), Int32(c)])
-        let conv = conv2DBackpropInput(
+        let conv = transposedConv2D(
             input.expandingShape(at: 1),
             shape: newShape,
             filter: filter.expandingShape(at: 0),
@@ -541,7 +541,7 @@ public struct TransposedConv2D<Scalar: TensorFlowFloatingPoint>: Layer {
           strides.1 + (filter.shape[1] * paddingIndex)
         let c = filter.shape[2]
         let newShape = Tensor<Int32>([Int32(batchSize), Int32(h), Int32(w), Int32(c)])
-        let conv = conv2DBackpropInput(
+        let conv = transposedConv2D(
             input,
             shape: newShape,
             filter: filter,

Sources/TensorFlow/Operators/NN.swift

@@ -149,6 +149,32 @@ func _vjpConv2D<Scalar: TensorFlowFloatingPoint>(
     })
 }
 
+/// Returns a 2-D transposed convolution with the specified input, filter, strides, and padding.
+///
+/// - Parameters:
+///   - input: The input.
+///   - shape: The output shape of the deconvolution operation.
+///   - filter: The convolution filter.
+///   - strides: The strides of the sliding filter for each dimension of the input.
+///   - padding: The padding for the operation
+///   - dilations: The dilation factor for each dimension of the input.
+/// - Precondition: `input` must have rank `4`.
+/// - Precondition: `filter` must have rank 4.
+@differentiable(wrt: (input, filter))
+public func transposedConv2D<Scalar: TensorFlowFloatingPoint>(
+    _ input: Tensor<Scalar>,
+    shape: Tensor<Int32>,
+    filter: Tensor<Scalar>,
+    strides: (Int, Int, Int, Int) = (1, 1, 1, 1),
+    padding: Padding = .valid,
+    dilations: (Int, Int, Int, Int) = (1, 1, 1, 1)
+) -> Tensor<Scalar> {
+    precondition(input.shape.rank == 4, "The input must have rank 4.")
+    precondition(filter.shape.rank == 4, "The filter must have rank 4.")
+    return conv2DBackpropInput(input, shape: shape, filter: filter,
+                               strides: strides, padding: padding, dilations: dilations)
+}
+
 /// TensorFlow builtin conv2d gradient helper for the input.
 @differentiable(wrt: (x, filter))
 @usableFromInline
@@ -170,8 +196,8 @@ func conv2DBackpropInput<Scalar: TensorFlowFloatingPoint>(
         dilations: [Int32(dilations.0), Int32(dilations.1), Int32(dilations.2), Int32(dilations.3)])
 }
 
-@usableFromInline
 @derivative(of: conv2DBackpropInput)
+@usableFromInline
 func _vjpConv2DBackpropInput<Scalar: TensorFlowFloatingPoint>(
     _ x: Tensor<Scalar>,
     _ shape: Tensor<Int32>,

Tests/TensorFlowTests/LayerTests.swift

@@ -360,8 +360,51 @@ final class LayerTests: XCTestCase {
         let expectedNoBias = Tensor<Float>(shape: [1, 4, 2, 1],
                                            scalars: [0, 4, 4, 20, 16, 56, 40, 104])
         XCTAssertEqual(outputNoBias, expectedNoBias)
-    }  
-
+    }
+    
+    func testTransposedConv2DGradient() {
+        let filter =  Tensor(shape: [3, 3, 2, 4], scalars: (0..<72).map(Float.init))
+        let bias = Tensor<Float>(zeros: [2])
+        let layer = TransposedConv2D<Float>(filter: filter,
+                                            bias: bias,
+                                            activation: identity,
+                                            strides: (2, 2),
+                                            padding: .same)
+        let input = Tensor(shape: [2, 2, 2, 4], scalars: (0..<32).map(Float.init))
+        let grads = gradient( at: input, layer) { $1($0).sum() }
+        // The expected value of the gradient was computed using the following Python code:
+        // ```
+        // import tensorflow as tf
+        // x = tf.reshape(tf.range(32, dtype=tf.float32), [2, 2, 2, 4])
+        // filter = tf.reshape(tf.range(72, dtype=tf.float32), [3, 3, 2, 4])
+        // bias = tf.zeros([2])
+        // with tf.GradientTape() as tape:
+        //     tape.watch([x, filter, bias])
+        //     y = tf.math.reduce_sum(tf.nn.conv2d_transpose(input=x,
+        //                                                   filters=filter,
+        //                                                   output_shape=[2, 4, 4, 2],
+        //                                                   strides=[1, 2, 2, 1],
+        //                                                   data_format="NHWC",
+        //                                                   padding="SAME") + bias)
+        // print(tape.gradient(y, [x, filter, bias]))
+        // ```
+        XCTAssertEqual(grads.0,
+                       [[[[612, 630, 648, 666], [360, 372, 384, 396]],
+                         [[264, 276, 288, 300], [144, 152, 160, 168]]],
+                        [[[612, 630, 648, 666], [360, 372, 384, 396]],
+                         [[264, 276, 288, 300], [144, 152, 160, 168]]]])
+        XCTAssertEqual(grads.1.filter,
+                       [[[[112, 120, 128, 136], [112, 120, 128, 136]],
+                         [[112, 120, 128, 136], [112, 120, 128, 136]],
+                         [[ 48,  52,  56,  60], [ 48,  52,  56,  60]]],
+                        [[[112, 120, 128, 136], [112, 120, 128, 136]],
+                         [[112, 120, 128, 136], [112, 120, 128, 136]],
+                         [[ 48,  52,  56,  60], [ 48,  52,  56,  60]]],
+                        [[[ 40,  44,  48,  52], [ 40,  44,  48,  52]],
+                         [[ 40,  44,  48,  52], [ 40,  44,  48,  52]],
+                         [[ 16,  18,  20,  22], [ 16,  18,  20,  22]]]])
+        XCTAssertEqual(grads.1.bias, [32, 32])
+    }
 
     func testTransposedConv3D() {
         let filter =  Tensor(shape: [2, 2, 2, 1, 1], scalars: (0..<8).map(Float.init))
@@ -1615,6 +1658,7 @@ final class LayerTests: XCTestCase {
         ("testConv3DGradient", testConv3DGradient),
         ("testTransposedConv1D", testTransposedConv1D),
         ("testTransposedConv2D", testTransposedConv2D),
+        ("testTransposedConv2DGradient", testTransposedConv2DGradient),
         ("testTransposedConv3D", testTransposedConv3D),
         ("testDepthwiseConv2D", testDepthwiseConv2D),
         ("testDepthwiseConv2DGradient", testDepthwiseConv2DGradient),