trim input

Author: RandyShuai

tests/test_backend.py

@@ -38,6 +38,8 @@
 _INPUT2 = "input2:0"
 _TFINPUT3 = "input3"
 _INPUT3 = "input3:0"
+_TFINPUT4 = "input4"
+_INPUT4 = "input4:0"
 _TFOUTPUT = "output"
 _OUTPUT = "output:0"
 _TFOUTPUT1 = "output1"
@@ -109,7 +111,7 @@ def _run_test_case(self, output_names_with_port, feed_dict, **kwargs):
         kwargs["convert_var_to_const"] = False
         kwargs["constant_fold"] = False
         return self.run_test_case(feed_dict, [], output_names_with_port, **kwargs)
-
+    '''
     def _test_expand_dims_known_rank(self, idx):
         tf.reset_default_graph()
         x_val = make_xval([3, 4])
@@ -2900,6 +2902,25 @@ def test_unique(self):
         # FIXME: indices in onnx are not the same as in tensorflow so don't check for now
         # self._run_test_case([_OUTPUT, _OUTPUT1], {_INPUT: x_val})
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
+    '''
+
+    def test_Conv2DBackpropInput_const(self):
+        input_sizes_val = np.array([1, 10, 10, 3], dtype=np.int32)
+        filter_val = np.random.randint(low=0, high=256, size=[3, 3, 3, 5]).astype(np.float32)
+        out_backprop_val = np.random.randint(low=0, high=256, size=[1, 10, 10, 5]).astype(np.float32)
+        _ = tf.nn.conv2d_backprop_input(input_sizes=input_sizes_val, filter=filter_val, out_backprop=out_backprop_val, strides=[1,1,1,1], padding='SAME', name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {})
+
+    def test_Conv2DBackpropInput(self):
+        input_sizes_val = np.array([1, 10, 10, 3], dtype=np.int32)
+        input_sizes = tf.placeholder(tf.int32, input_sizes_val.shape, name=_TFINPUT)
+        filter_val = np.random.randint(low=0, high=256, size=[3, 3, 3, 5]).astype(np.float32)
+        filter = tf.placeholder(tf.float32, filter_val.shape, name=_TFINPUT1)
+        out_backprop_val = np.random.randint(low=0, high=256, size=[1, 10, 10, 5]).astype(np.float32)
+        out_backprop = tf.placeholder(tf.float32, out_backprop_val.shape, name=_TFINPUT2)
+        _ = tf.nn.conv2d_backprop_input(input_sizes, filter, out_backprop, strides=[1,1,1,1], padding='SAME', name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: input_sizes_val, _INPUT1: filter_val, _INPUT2: out_backprop_val})
+
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/onnx_opset/nn.py

@@ -226,21 +226,43 @@ def version_1(cls, ctx, node, **kwargs):
         # Note: inputs are reversed from what one would expect.
         kernel_shape = conv_kernel_shape(ctx, node, 1)
         input_shape = ctx.get_shape(node.input[2])
+        append_slice = False
 
         # ouput_shape is explicitly specified here, in this case pads values are auto generated/calculated.
-        output_shape = ctx.get_shape(node.output[0])
-        if node.is_nhwc():
-            new_output_shape = [output_shape[1], output_shape[2]]
-            input_hw = [input_shape[1], input_shape[2]]
+        if node.inputs[0].is_const():
+            output_shape = ctx.get_shape(node.output[0])
+            if node.is_nhwc():
+                new_output_shape = [output_shape[1], output_shape[2]]
+                input_hw = [input_shape[1], input_shape[2]]
+            else:
+                new_output_shape = [output_shape[2], output_shape[3]]
+                input_hw = [input_shape[2], input_shape[3]]
+            utils.make_sure(new_output_shape.count(-1) <= 0, "output h and w need to be known")
+            utils.make_sure(new_output_shape[0] >= input_hw[0] and new_output_shape[1] >= input_hw[1],
+                            "output h and w cannot be smaller than input h and w.")
+            node.set_attr("output_shape", new_output_shape)
         else:
-            new_output_shape = [output_shape[2], output_shape[3]]
-            input_hw = [input_shape[2], input_shape[3]]
-
-        utils.make_sure(new_output_shape.count(-1) <= 0, "output h and w need to be known")
-        utils.make_sure(new_output_shape[0] >= input_hw[0] and new_output_shape[1] >= input_hw[1],
-                        "output h and w cannot be smaller than input h and w.")
-
-        node.set_attr("output_shape", new_output_shape)
+            input_shape = ctx.make_node("Cast", [node.input[0]], attr={'to': TensorProto.INT64})
+            output_shape = ctx.make_node("Shape", [node.output[0]])
+            output_h = GraphBuilder(ctx).make_slice({"data": output_shape.output[0], "ends": [2], "starts": [1], "axes": [0]})
+            output_w = GraphBuilder(ctx).make_slice({"data": output_shape.output[0], "ends": [3], "starts": [2], "axes": [0]})
+            expect_h = GraphBuilder(ctx).make_slice({"data": input_shape.output[0], "ends": [2], "starts": [1], "axes": [0]})
+            expect_w = GraphBuilder(ctx).make_slice({"data": input_shape.output[0], "ends": [3], "starts": [2], "axes": [0]})
+            diff_h = ctx.make_node("Sub", [output_h, expect_h])
+            diff_w = ctx.make_node("Sub", [output_w, expect_w])
+            const_two = ctx.make_const(utils.make_name(node.name + "_const_two"), np.array([2], dtype=np.int64))
+            start_h = ctx.make_node("Div", [diff_h.output[0], const_two.output[0]])
+            start_w = ctx.make_node("Div", [diff_w.output[0], const_two.output[0]])
+            end_h = ctx.make_node("Add", [start_h.output[0], expect_h])
+            end_w = ctx.make_node("Add", [start_w.output[0], expect_w])
+            starts = ctx.make_node("Concat", [start_h.output[0], start_w.output[0]], attr={"axis":0})
+            ends = ctx.make_node("Concat", [end_h.output[0], end_w.output[0]], attr={"axis":0}, name="concat_efgh")
+            const_one_two = ctx.make_const(utils.make_name(node.name + "_const_one_two"), np.array([1,2], dtype=np.int64))
+            slice_node = ctx.make_node("Slice", [node.output[0], starts.output[0], ends.output[0], const_one_two.output[0]])
+            downstream_nodes = ctx.find_output_consumers(node.output[0])
+            downstream_nodes.remove(output_shape)
+            downstream_nodes.remove(slice_node)
+            ctx.replace_all_inputs(downstream_nodes, node.output[0], slice_node.output[0])
 
         strides = conv_dims_attr(node, "strides")
         conv_dims_attr(node, "dilations")