From 1239de3efec67db3de0974f86958973018743a3a Mon Sep 17 00:00:00 2001
From: Hariharan Seshadri <shariharan91@gmail.com>
Date: Wed, 29 Jan 2020 15:42:43 -0800
Subject: [PATCH] Fix output shape computation in ConvTranspose op (#2688)

* Fix output shape computation in ConvTranspose op

* Update test
---
 .../cpu/nn/conv_transpose_attributes.h        |  6 ++--
 .../cpu/nn/conv_transpose_op_test.cc          | 35 ++++++++++++++-----
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h b/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h
index a8cd1f5a9c..10e93e6a46 100644
--- a/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h
+++ b/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h
@@ -188,7 +188,7 @@ struct ConvTransposeAttributes : public ConvAttributes {
     if (*out_size != -1) {
       ORT_ENFORCE(*out_size >= 0);
       // total padding size
-      int64_t paddings = std::max<int64_t>(0, (in_size - 1) * stride + kernel + dilation - 1 + adj - *out_size);
+      int64_t paddings = std::max<int64_t>(0, (in_size - 1) * stride + adj + (kernel - 1) * dilation + 1 - *out_size);
       if (pad_type == AutoPadType::SAME_UPPER) {  // pad more on head when paddings are odd.
         *pad_head = paddings - paddings / 2;
         *pad_tail = paddings / 2;
@@ -210,14 +210,14 @@ struct ConvTransposeAttributes : public ConvAttributes {
         case AutoPadType::SAME_LOWER:
           *pad_head = 0;
           *pad_tail = 0;
-          *out_size = (in_size - 1) * stride + kernel + dilation - 1 + adj;
+          *out_size = (in_size - 1) * stride + adj + (kernel - 1) * dilation + 1;
           break;
         default:
           throw NotImplementedException("pad type not supported");
       }
     } else {
       *out_size =
-          (in_size - 1) * stride + kernel + dilation - 1 + adj - *pad_head - *pad_tail;
+          (in_size - 1) * stride + adj + (kernel - 1) * dilation + 1 - *pad_head - *pad_tail;
     }
   }
 };
diff --git a/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc b/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc
index 7be5cadc39..02c0835757 100644
--- a/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc
+++ b/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc
@@ -54,20 +54,19 @@ void TestConvTransposeOp(const ConvTransposeOpAttributes& attributes,
   }
   test.AddOutput<float>("Y", expected_output_shape, expected_output);
 
-  
   test.Run(expect_result, err_str, excluded_provider_types);  // Disable TensorRT because weight as input is not supported
 }
 }  // namespace
 
 TEST(ConvTransposeTest, ConvTranspose_1D) {
   ConvTransposeOpAttributes attrs = {
-      vector<int64_t>{3},        // kernel_shape
-      {},                        // output_padding
-      {},                        // output_shape
-      vector<int64_t>{0, 0},     // pads
-      vector<int64_t>{1},        // strides
-      vector<int64_t>{1},        // dilations
-      1                          // group
+      vector<int64_t>{3},     // kernel_shape
+      {},                     // output_padding
+      {},                     // output_shape
+      vector<int64_t>{0, 0},  // pads
+      vector<int64_t>{1},     // strides
+      vector<int64_t>{1},     // dilations
+      1                       // group
   };
   vector<float> X = {0.0f, 1.0f, 2.0f};
   vector<int64_t> X_shape = {1, 1, 3};
@@ -533,5 +532,25 @@ TEST(ConvTransposeTest, ConvTranspose_DefaultStridesAndDilations) {
   TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape);
 }
 
+TEST(ConvTransposeTest, ConvTranspose_2D_NonDefaultStridesAndDilations) {
+  ConvTransposeOpAttributes attrs = {
+      vector<int64_t>{1, 4},        // kernel_shape
+      {},                           // output_padding
+      {},                           // output_shape
+      vector<int64_t>{0, 0, 0, 0},  // pads
+      vector<int64_t>{1, 2},        // strides
+      vector<int64_t>{1, 3},        // dilations
+      1                             // group
+  };
+  vector<float> X = {1., 2.};
+  vector<int64_t> X_shape = {1, 1, 1, 2};
+  vector<float> W = {1., 1., 1., 1.};
+  vector<int64_t> W_shape = {1, 1, 1, 4};
+  vector<int64_t> Y_shape = {1, 1, 1, 12};
+  auto expected_vals = {1.f, 0.f, 2.f, 1.f, 0.f, 2.f, 1.f, 0.f, 2.f, 1.f, 0.f, 2.f};
+
+  TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape);
+}
+
 }  // namespace test
 }  // namespace onnxruntime