Fix implementation issue of convTranspose2d to support options.groups

src/conv_transpose2d.js

@@ -1,7 +1,7 @@
 'use strict';
 
 import {Tensor} from './lib/tensor.js';
-import {validateConv2dParams} from './lib/validate-input.js';
+import {validateConvTranspose2dParams} from './lib/validate-input.js';
 import {computePaddingForAutoPad} from './lib/compute-padding.js';
 import {transpose} from './transpose.js';
 
@@ -44,10 +44,10 @@ export function convTranspose2d(
     filter = transpose(filter, {permutation: [0, 3, 1, 2]});
   }
 
-  validateConv2dParams(input, filter, {groups, bias});
+  validateConvTranspose2dParams(input, filter, {groups, bias});
 
   const [batchCount, inputChannels, inputHeight, inputWidth] = input.shape;
-  const [outputChannels, , filterHeight, filterWidth] = filter.shape;
+  const [outputChannelsPerGroup, , filterHeight, filterWidth] = filter.shape;
   const [strideHeight, strideWidth] = strides;
   const [dilationHeight, dilationWidth] = dilations;
   const effectiveFilterHeight = (filterHeight - 1) * dilationHeight + 1;
@@ -76,6 +76,8 @@ export function convTranspose2d(
   let outputHeight;
   let outputWidth;
   outputShape[0] = batchCount;
+
+  const outputChannels = outputChannelsPerGroup * groups;
   outputShape[1] = outputChannels;
 
   if (outputSizes === undefined) {
@@ -104,7 +106,6 @@ export function convTranspose2d(
     dilationWidth * (filter.shape[3] - 1) - beginningPaddingWidth;
   const [realStrideHeight, realStrideWidth] = [1, 1];
 
-  const outputChannelsPerGroup = outputChannels / groups;
   const inputChannelsPerGroup = inputChannels / groups;
 
   for (let ib = 0; ib < batchCount; ++ib) {
@@ -115,8 +116,8 @@ export function convTranspose2d(
             ih += realStrideHeight, ++oh) {
             for (let iw = -realBeginningPaddingWidth, ow = 0; ow < outputWidth;
               iw += realStrideWidth, ++ow) {
-              const effectiveOutputChannel = oc + g * outputChannelsPerGroup;
-              const outputLocation = [ib, effectiveOutputChannel, oh, ow];
+              const effectiveOutputChannel = oc;
+              const outputLocation = [ib, oc + g * outputChannelsPerGroup, oh, ow];
               // make filter rotate 180
               for (let kh = filterHeight + (filterHeight - 1) * (dilationHeight - 1) - 1;
                 kh >= 0; --kh) {

src/lib/validate-input.js

@@ -99,6 +99,28 @@ export function validateConv2dParams(input, filter, {bias, groups = 1}) {
   }
 }
 
+export function validateConvTranspose2dParams(input, filter, {bias, groups = 1}) {
+  const inputChannels = input.shape[1];
+  // filter of oihw
+  const outputChannelsPerGroup = filter.shape[0];
+
+  if (input.rank !== 4) {
+    throw new Error('The input should be a 4-D tensor.');
+  }
+  if (filter.rank !== 4) {
+    throw new Error('The filter should be a 4-D tensor.');
+  }
+  if (inputChannels % groups !== 0) {
+    throw new Error('The input channels is invalid.');
+  }
+  if (inputChannels !== filter.shape[1]) {
+    throw new Error('The input channels of filter is invalid.');
+  }
+  if (bias && (bias.rank !== 1 || bias.shape[0] != outputChannelsPerGroup * groups)) {
+    throw new Error('the bias should be a 1-D tensor with the shape of [output_channels].');
+  }
+}
+
 export function validateGemmParams(a, b) {
   if (a.rank !== 2) {
     throw new Error('The input a is not a 2-D tensor.');

test/conv_transpose2d_test.js

@@ -55,6 +55,81 @@ describe('test convTranspose2d', function() {
     testConvTranspose2d(input, filter, expected);
   });
 
+  it('convTranspose2d options.groups', function() {
+    const input = {
+      shape: [1, 2, 2, 2],
+      data: [
+        2,  4,
+        0,  1,
+        2,  4,
+        0,  1,
+      ],
+    };
+    const filter = {
+      shape: [2, 1, 2, 2],
+      data: [
+        3,  1,
+        1,  5,
+        3,  1,
+        1,  5,
+      ],
+    };
+    const options = {
+      groups: 2,
+    };
+    const expected = {
+      shape: [1, 2, 3, 3],
+      data: [
+        6, 14,  4,
+        2, 17, 21,
+        0,  1,  5,
+        6, 14,  4,
+        2, 17, 21,
+        0,  1,  5,
+      ],
+    };
+    testConvTranspose2d(input, filter, expected, options);
+  });
+
+  it('convTranspose2d options.groups=2 options.strides=[2, 2]', function() {
+    const input = {
+      shape: [1, 2, 2, 2],
+      data: [
+        2,  4,
+        0,  1,
+        2,  4,
+        0,  1,
+      ],
+    };
+    const filter = {
+      shape: [2, 1, 2, 2],
+      data: [
+        3,  1,
+        1,  5,
+        3,  1,
+        1,  5,
+      ],
+    };
+    const options = {
+      groups: 2,
+      strides: [2, 2],
+    };
+    const expected = {
+      shape: [1, 2, 4, 4],
+      data: [
+        6,  2, 12,  4,
+        2, 10,  4, 20,
+        0,  0,  3,  1,
+        0,  0,  1,  5,
+        6,  2, 12,  4,
+        2, 10,  4, 20,
+        0,  0,  3,  1,
+        0,  0,  1,  5,
+      ],
+    };
+    testConvTranspose2d(input, filter, expected, options);
+  });
+
   it('convTranspose2d options.padding', function() {
     const input = {
       shape: [1, 1, 2, 2],