Nd deconv read support (#554)

* initial progress

* formatting

* check existing tests

* formatting

* change for loop to transform

* formatting

* add tests

* formatting

* remove comment

* add more tests

* remove extra slice axes
Co-authored-by: mvermeulen <5479696+mvermeulen@users.noreply.github.com>

src/include/migraphx/op/deconvolution.hpp

@@ -19,9 +19,9 @@ namespace op {
 
 struct deconvolution
 {
-    std::array<std::size_t, 2> padding  = {{0, 0}};
-    std::array<std::size_t, 2> stride   = {{1, 1}};
-    std::array<std::size_t, 2> dilation = {{1, 1}};
+    std::vector<std::size_t> padding  = {0, 0};
+    std::vector<std::size_t> stride   = {1, 1};
+    std::vector<std::size_t> dilation = {1, 1};
 
     padding_mode_t padding_mode = default_;
     int group                   = 1;
@@ -39,25 +39,27 @@ struct deconvolution
     std::string name() const { return "deconvolution"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(2).same_type().same_ndims().only_dims(4);
+        check_shapes{inputs, *this}.has(2).same_type().same_ndims().min_ndims(3);
+        if(not(padding.size() == stride.size() and padding.size() == dilation.size()))
+        {
+            MIGRAPHX_THROW("deconvolution: inconsistent attribute sizes");
+        }
 
         const shape& input   = inputs.at(0);
         const shape& weights = inputs.at(1);
         auto t               = input.type();
+        size_t kdims         = input.lens().size() - 2;
+
+        std::vector<size_t> output_lens{input.lens()[0], weights.lens()[1]};
 
-        return {t,
+        for(size_t i = 0; i < kdims; i++)
         {
-                    input.lens()[0],
-                    weights.lens()[1],
-                    std::size_t(std::max<std::ptrdiff_t>(
+            output_lens.push_back(std::size_t(std::max<std::ptrdiff_t>(
                 1,
-                        stride[0] * (input.lens()[2] - 1) +
-                            ((weights.lens()[2] - 1) * dilation[0] + 1) - 2 * padding[0])),
-                    std::size_t(std::max<std::ptrdiff_t>(
-                        1,
-                        stride[1] * (input.lens()[3] - 1) +
-                            ((weights.lens()[3] - 1) * dilation[1] + 1) - 2 * padding[1])),
-                }};
+                stride[i] * (input.lens()[i + 2] - 1) +
+                    ((weights.lens()[i + 2] - 1) * dilation[i] + 1) - 2 * padding[i])));
+        }
+        return {t, output_lens};
     }
 };
 

src/onnx/onnx.cpp

@@ -320,29 +320,32 @@ struct onnx_parser
         return curr_ins;
     }
 
-    template <class Op>
-    void check_asym_padding(instruction_ref& ins,
-                            const std::vector<int64_t>& padding,
-                            Op& op,
-                            float pad_val = 0)
+    bool is_asym_padding(const std::vector<int64_t>& padding)
     {
-        bool asym_padding = false;
         assert(padding.size() % 2 == 0);
         size_t pad_ndims = padding.size() / 2;
 
-        auto left_pad_it  = padding.begin();
-        auto right_pad_it = left_pad_it + pad_ndims;
-
         for(size_t i = 0; i < pad_ndims; i++)
         {
             if(padding[i] != padding[i + pad_ndims])
             {
-                asym_padding = true;
-                break;
+                return true;
+            }
         }
+        return false;
     }
 
-        if(asym_padding)
+    template <class Op>
+    void check_asym_padding(instruction_ref& ins,
+                            const std::vector<int64_t>& padding,
+                            Op& op,
+                            float pad_val = 0)
+    {
+        size_t pad_ndims  = padding.size() / 2;
+        auto left_pad_it  = padding.begin();
+        auto right_pad_it = left_pad_it + pad_ndims;
+
+        if(is_asym_padding(padding))
         {
             std::vector<int64_t> asym_pads{0, 0, 0, 0}; // don't pad N and C
             // add left pads
@@ -654,7 +657,11 @@ struct onnx_parser
         op::deconvolution op;
         auto l0 = args[0];
         std::vector<std::int64_t> padding;
-        bool asymm_padding = false;
+        bool asym_padding = false;
+        auto in_lens      = l0->get_shape().lens();
+        assert(in_lens.size() > 2);
+        auto kdims = in_lens.size() - 2;
+
         if(contains(info.attributes, "pads"))
         {
             if(contains(info.attributes, "auto_pad"))
@@ -662,38 +669,45 @@ struct onnx_parser
                 auto s = info.attributes["auto_pad"].s();
                 if(contains(info.attributes, "pads") and to_upper(s) != "NOTSET")
                 {
-                    MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
+                    MIGRAPHX_THROW("PARSE_CONV_TRANSPOSE: auto_pad and padding cannot be specified "
+                                   "simultaneously");
                 }
             }
             copy(info.attributes["pads"].ints(), std::back_inserter(padding));
-            if(padding.size() != 4)
-            {
-                MIGRAPHX_THROW("padding should have 4 values");
-            }
-            if(padding[0] != padding[2] || padding[1] != padding[3])
-            {
-                asymm_padding = true;
-            }
-            else
+
+            asym_padding = is_asym_padding(padding);
+
+            if(not asym_padding)
             {
-                op.padding[0] = padding[0];
-                op.padding[1] = padding[1];
+                size_t pad_ndims = padding.size() / 2;
+                check_attr_sizes(kdims, pad_ndims, "PARSE_CONV_TRANSPOSE: inconsistent paddings");
+                op.padding.clear();
+                std::transform(padding.begin(),
+                               padding.begin() + pad_ndims,
+                               std::back_inserter(op.padding),
+                               [](auto pad_val) { return pad_val; });
             }
         }
         if(contains(info.attributes, "strides"))
         {
-            copy(info.attributes["strides"].ints(), op.stride.begin());
+            op.stride.clear();
+            copy(info.attributes["strides"].ints(), std::back_inserter(op.stride));
+            check_attr_sizes(kdims, op.stride.size(), "PARSE_CONV_TRANSPOSE: inconsistent strides");
         }
         if(contains(info.attributes, "dilations"))
         {
-            copy(info.attributes["dilations"].ints(), op.dilation.begin());
+            op.dilation.clear();
+            copy(info.attributes["dilations"].ints(), std::back_inserter(op.dilation));
+            check_attr_sizes(
+                kdims, op.dilation.size(), "PARSE_CONV_TRANSPOSE: inconsistent dilations");
         }
         if(contains(info.attributes, "auto_pad"))
         {
             auto s = info.attributes["auto_pad"].s();
             if(contains(info.attributes, "pads") and to_upper(s) != "NOTSET")
             {
-                MIGRAPHX_THROW("auto_pad and padding cannot be specified simultaneously");
+                MIGRAPHX_THROW("PARSE_CONV_TRANSPOSE: auto_pad and padding cannot be specified "
+                               "simultaneously");
             }
 
             if(s.find("SAME") != std::string::npos)
@@ -707,25 +721,37 @@ struct onnx_parser
             op.group = parse_value(info.attributes.at("group")).at<int>();
         }
 
+        recalc_conv_attributes(op, kdims);
+
         auto l1                   = prog.add_instruction(op, l0, args[1]);
         std::vector<int64_t> dims = to_int64_vector(l1->get_shape().lens());
-        std::vector<int64_t> curr_shape{dims[2], dims[3]};
-        if(asymm_padding)
+        std::vector<int64_t> curr_shape(dims.begin() + 2, dims.end());
+        if(asym_padding)
         {
-            op::slice slice_op;
-            slice_op.axes   = {0, 1, 2, 3};
-            slice_op.starts = {0, 0, 0 + padding[0], 0 + padding[1]};
-            slice_op.ends   = {
-                dims[0], dims[1], curr_shape[0] - padding[2], curr_shape[1] - padding[3]};
+            std::vector<int64_t> axes(kdims);
+            std::iota(axes.begin(), axes.end(), 2); // ignore first 2 dims
+
+            auto pad_kdim_start = padding.begin() + kdims;
+            std::vector<int64_t> starts(padding.begin(), pad_kdim_start);
+
+            std::vector<int64_t> ends{};
+            std::transform(curr_shape.begin(),
+                           curr_shape.end(),
+                           pad_kdim_start,
+                           std::back_inserter(ends),
+                           [](auto curr_dim, auto pad_dim) { return curr_dim - pad_dim; });
 
-            l1 = prog.add_instruction(slice_op, l1);
+            l1 = prog.add_instruction(op::slice{axes, starts, ends}, l1);
         }
 
         if(contains(info.attributes, "output_padding"))
         {
-            std::vector<int64_t> output_padding;
+            size_t non_kdims = dims.size() * 2 - kdims;
+            std::vector<int64_t> output_padding(non_kdims, 0);
             copy(info.attributes["output_padding"].ints(), std::back_inserter(output_padding));
-            output_padding = {0, 0, 0, 0, 0, 0, output_padding[0], output_padding[1]};
+            check_attr_sizes(kdims,
+                             output_padding.size() - non_kdims,
+                             "PARSE_CONV_TRANSPOSE: inconsistent output padding");
             l1 = prog.add_instruction(op::pad{output_padding}, l1);
         }
 
@@ -733,18 +759,18 @@ struct onnx_parser
         {
             std::vector<int64_t> output_shape;
             copy(info.attributes["output_shape"].ints(), std::back_inserter(output_shape));
+            check_attr_sizes(
+                kdims, output_shape.size(), "PARSE_CONV_TRANSPOSE: inconsistent output shape");
             dims = to_int64_vector(l1->get_shape().lens());
-            curr_shape = {dims[2], dims[3]};
+            copy(dims.begin() + 2, dims.end(), curr_shape.begin());
             if(curr_shape != output_shape)
             {
-                std::vector<int64_t> target_padding = {0,
-                                                       0,
-                                                       0,
-                                                       0,
-                                                       0,
-                                                       0,
-                                                       output_shape[0] - curr_shape[0],
-                                                       output_shape[1] - curr_shape[1]};
+                std::vector<int64_t> target_padding(dims.size() * 2 - kdims, 0);
+                std::transform(output_shape.begin(),
+                               output_shape.end(),
+                               curr_shape.begin(),
+                               std::back_inserter(target_padding),
+                               [](auto out_dim, auto curr_dim) { return out_dim - curr_dim; });
                 l1 = prog.add_instruction(op::pad{target_padding}, l1);
             }
         }

test/onnx/deconv_output_padding_3d_test.onnx

+deconv_output_padding_3d_test:

+G
+
x
+
w
y"
ConvTranspose*
+output_padding@
@
@

*
+strides@@@
deconv_output_padding_3d_testZ
+
x
+

+

+

+
+
+Z
+
w
+

+

+
+
+
+b
+
y
+

+

+
+
+
+
+B
\ No newline at end of file

test/onnx/deconv_output_shape_3d_test.onnx

+deconv_output_shape_3d_test:

+E
+
x
+
w
y"
ConvTranspose*
+output_shape@
+@@
*
+strides@@@
deconv_output_shape_3d_testZ
+
x
+

+

+

+
+
+Z
+
w
+

+

+
+
+
+b
+
y
+

+

+
+
+
+
+B
\ No newline at end of file

test/onnx/gen_onnx.py

@@ -871,7 +871,23 @@ def deconv_input_pads_asymm_test():
 
 
 @onnx_test
-def deconv_output_shape_test():
+def deconv_input_pads_asymm_1d_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 1, 3])
+    w = helper.make_tensor_value_info('w', TensorProto.FLOAT, [1, 2, 3])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 2, 6])
+
+    node = onnx.helper.make_node('ConvTranspose',
+                                 inputs=['x', 'w'],
+                                 outputs=['y'],
+                                 strides=[2],
+                                 pads=[0, 1],
+                                 dilations=[1])
+
+    return ([node], [x, w], [y])
+
+
+@onnx_test
+def deconv_output_padding_test():
     x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 1, 3, 3])
     w = helper.make_tensor_value_info('w', TensorProto.FLOAT, [1, 2, 3, 3])
     y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 2, 10, 8])
@@ -880,13 +896,28 @@ def deconv_output_shape_test():
                                  inputs=['x', 'w'],
                                  outputs=['y'],
                                  strides=[3, 2],
-                                 output_shape=[10, 8])
+                                 output_padding=[1, 1])
 
     return ([node], [x, w], [y])
 
 
 @onnx_test
-def deconv_output_padding_test():
+def deconv_output_padding_3d_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 1, 3, 3, 3])
+    w = helper.make_tensor_value_info('w', TensorProto.FLOAT, [1, 2, 3, 3, 3])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 2, 10, 8, 8])
+
+    node = onnx.helper.make_node('ConvTranspose',
+                                 inputs=['x', 'w'],
+                                 outputs=['y'],
+                                 strides=[3, 2, 2],
+                                 output_padding=[1, 1, 1])
+
+    return ([node], [x, w], [y])
+
+
+@onnx_test
+def deconv_output_shape_test():
     x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 1, 3, 3])
     w = helper.make_tensor_value_info('w', TensorProto.FLOAT, [1, 2, 3, 3])
     y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 2, 10, 8])
@@ -895,7 +926,22 @@ def deconv_output_padding_test():
                                  inputs=['x', 'w'],
                                  outputs=['y'],
                                  strides=[3, 2],
-                                 output_padding=[1, 1])
+                                 output_shape=[10, 8])
+
+    return ([node], [x, w], [y])
+
+
+@onnx_test
+def deconv_output_shape_3d_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 1, 3, 3, 3])
+    w = helper.make_tensor_value_info('w', TensorProto.FLOAT, [1, 2, 3, 3, 3])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 2, 10, 8, 8])
+
+    node = onnx.helper.make_node('ConvTranspose',
+                                 inputs=['x', 'w'],
+                                 outputs=['y'],
+                                 strides=[3, 2, 2],
+                                 output_shape=[10, 8, 8])
 
     return ([node], [x, w], [y])
 

test/onnx/onnx_test.cpp

@@ -633,13 +633,25 @@ TEST_CASE(deconv_input_pads_asymm_test)
     auto l0 = p.add_parameter("x", {migraphx::shape::float_type, {1, 1, 3, 3}});
     auto l1 = p.add_parameter("w", {migraphx::shape::float_type, {1, 2, 3, 3}});
     auto l2 = p.add_instruction(migraphx::op::deconvolution{{0, 0}, {3, 2}}, l0, l1);
-    p.add_instruction(migraphx::op::slice{{0, 1, 2, 3}, {0, 0, 0, 0}, {1, 2, 8, 6}}, l2);
+    p.add_instruction(migraphx::op::slice{{2, 3}, {0, 0}, {8, 6}}, l2);
 
     auto prog = optimize_onnx("deconv_input_pads_asymm_test.onnx");
     EXPECT(p == prog);
 }
 
-TEST_CASE(deconv_output_shape_test)
+TEST_CASE(deconv_input_pads_asymm_1d_test)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("x", {migraphx::shape::float_type, {1, 1, 3}});
+    auto l1 = p.add_parameter("w", {migraphx::shape::float_type, {1, 2, 3}});
+    auto l2 = p.add_instruction(migraphx::op::deconvolution{{0}, {2}, {1}}, l0, l1);
+    p.add_instruction(migraphx::op::slice{{2}, {0}, {6}}, l2);
+
+    auto prog = optimize_onnx("deconv_input_pads_asymm_1d_test.onnx");
+    EXPECT(p == prog);
+}
+
+TEST_CASE(deconv_output_padding_test)
 {
     migraphx::program p;
     auto l0 = p.add_parameter("x", {migraphx::shape::float_type, {1, 1, 3, 3}});
@@ -647,11 +659,24 @@ TEST_CASE(deconv_output_shape_test)
     auto l2 = p.add_instruction(migraphx::op::deconvolution{{0, 0}, {3, 2}}, l0, l1);
     p.add_instruction(migraphx::op::pad{{0, 0, 0, 0, 0, 0, 1, 1}}, l2);
 
-    auto prog = optimize_onnx("deconv_output_shape_test.onnx");
+    auto prog = optimize_onnx("deconv_output_padding_test.onnx");
     EXPECT(p == prog);
 }
 
-TEST_CASE(deconv_output_padding_test)
+TEST_CASE(deconv_output_padding_3d_test)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("x", {migraphx::shape::float_type, {1, 1, 3, 3, 3}});
+    auto l1 = p.add_parameter("w", {migraphx::shape::float_type, {1, 2, 3, 3, 3}});
+    auto l2 =
+        p.add_instruction(migraphx::op::deconvolution{{0, 0, 0}, {3, 2, 2}, {1, 1, 1}}, l0, l1);
+    p.add_instruction(migraphx::op::pad{{0, 0, 0, 0, 0, 0, 0, 1, 1, 1}}, l2);
+
+    auto prog = optimize_onnx("deconv_output_padding_3d_test.onnx");
+    EXPECT(p == prog);
+}
+
+TEST_CASE(deconv_output_shape_test)
 {
     migraphx::program p;
     auto l0 = p.add_parameter("x", {migraphx::shape::float_type, {1, 1, 3, 3}});
@@ -659,7 +684,20 @@ TEST_CASE(deconv_output_padding_test)
     auto l2 = p.add_instruction(migraphx::op::deconvolution{{0, 0}, {3, 2}}, l0, l1);
     p.add_instruction(migraphx::op::pad{{0, 0, 0, 0, 0, 0, 1, 1}}, l2);
 
-    auto prog = optimize_onnx("deconv_output_padding_test.onnx");
+    auto prog = optimize_onnx("deconv_output_shape_test.onnx");
+    EXPECT(p == prog);
+}
+
+TEST_CASE(deconv_output_shape_3d_test)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("x", {migraphx::shape::float_type, {1, 1, 3, 3, 3}});
+    auto l1 = p.add_parameter("w", {migraphx::shape::float_type, {1, 2, 3, 3, 3}});
+    auto l2 =
+        p.add_instruction(migraphx::op::deconvolution{{0, 0, 0}, {3, 2, 2}, {1, 1, 1}}, l0, l1);
+    p.add_instruction(migraphx::op::pad{{0, 0, 0, 0, 0, 0, 0, 1, 1, 1}}, l2);
+
+    auto prog = optimize_onnx("deconv_output_shape_3d_test.onnx");
     EXPECT(p == prog);
 }
 

test/op_shape_test.cpp

@@ -89,6 +89,28 @@ TEST_CASE(convolution_shape)
                  weights_3d);
 }
 
+TEST_CASE(deconvolution_shape)
+{
+    migraphx::shape input{migraphx::shape::float_type, {4, 4, 1, 1}};
+    migraphx::shape output{migraphx::shape::float_type, {4, 3, 3, 3}};
+    migraphx::shape weights{migraphx::shape::float_type, {4, 3, 3, 3}};
+    expect_shape(output, migraphx::op::deconvolution{}, input, weights);
+    throws_shape(migraphx::op::deconvolution{}, input);
+
+    migraphx::shape input_1d{migraphx::shape::float_type, {4, 4, 1}};
+    migraphx::shape output_1d{migraphx::shape::float_type, {4, 3, 3}};
+    migraphx::shape weights_1d{migraphx::shape::float_type, {4, 3, 3}};
+    expect_shape(output_1d, migraphx::op::deconvolution{{0}, {1}, {1}}, input_1d, weights_1d);
+
+    migraphx::shape input_3d{migraphx::shape::float_type, {4, 4, 1, 1, 1}};
+    migraphx::shape output_3d{migraphx::shape::float_type, {4, 3, 3, 3, 3}};
+    migraphx::shape weights_3d{migraphx::shape::float_type, {4, 3, 3, 3, 3}};
+    expect_shape(output_3d,
+                 migraphx::op::deconvolution{{0, 0, 0}, {1, 1, 1}, {1, 1, 1}},
+                 input_3d,
+                 weights_3d);
+}
+
 TEST_CASE(quant_convolution_shape)
 {
     migraphx::shape output{migraphx::shape::int32_type, {4, 4, 1, 1}};
@@ -115,6 +137,7 @@ TEST_CASE(inconsistent_attr_shape)
     migraphx::shape input{migraphx::shape::float_type, {4, 3, 3, 3}};
     migraphx::shape weights{migraphx::shape::float_type, {4, 3, 3, 3}};
     throws_shape(migraphx::op::convolution{{1, 1}, {2}, {3, 3, 3}}, input, weights);
+    throws_shape(migraphx::op::deconvolution{{1, 1}, {2}, {3, 3, 3}}, input, weights);
     throws_shape(migraphx::op::pooling{"max", {1}, {0}, {1, 1}}, input);
 }