Merge pull request #1071 from NikolasMarkou/DepthwiseConv2d_translation

#1070 replacing Reshape and Transpose operators to the kernel that are  invalid in tensorrt with just a reshaped tensor

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -34,6 +34,7 @@
 
 _STRIDE1x1 = [1, 1, 1, 1]
 _KERNEL3x3 = [3, 3, 1, 1]
+_DILATIONS1x1 = [1, 1, 1, 1]
 
 # names for input and outputs for tests
 _TFINPUT = "input"
@@ -348,7 +349,7 @@ def _conv_test(self, x_val, w, strides=None, padding="VALID", dilations=None, rt
         if strides is None:
             strides = _STRIDE1x1
         if dilations is None:
-            dilations = [1, 1, 1, 1]
+            dilations = _DILATIONS1x1
         def func(x):
             kernel = tf.constant(w, dtype=tf.float32, name='k')
             conv = tf.nn.conv2d(x, kernel, strides=strides, padding=padding, dilations=dilations)
@@ -3580,6 +3581,27 @@ def func(y, x):
         self._run_test_case(
             func, [_OUTPUT], {_INPUT: y_val, _INPUT2: x_val}, rtol=1e-06)
 
+    def _conv_kernel_as_input_test(self, x_val, w_val, strides=None,
+                                   padding="VALID", dilations=None, rtol=1e-07):
+        if strides is None:
+            strides = _STRIDE1x1
+        if dilations is None:
+            dilations = _DILATIONS1x1
+
+        def func(x, kernel):
+            conv = tf.nn.conv2d(x, kernel, strides=strides, padding=padding,
+                                dilations=dilations)
+            return tf.identity(conv, name=_TFOUTPUT)
+
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT2: w_val}, rtol=rtol)
+
+    def test_conv2d_1_kernel_as_input(self):
+        x_val = make_xval((1, 1, 5, 5)).transpose(NCHW_TO_NHWC)
+        w_val = np.array([[2., 1., 1.],
+                          [1., 3., 1.],
+                          [1., 1., 4.]], dtype=np.float32).reshape(_KERNEL3x3)
+        self._conv_kernel_as_input_test(x_val, w_val)
+
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/onnx_opset/nn.py

@@ -117,24 +117,28 @@ def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
     if with_kernel:
         # Some ONNX convolution ops require to reshape the kernel (ie. depthwise_conv2d).
         if new_kernel_shape:
-            kernel_name = node.input[1]
-
-            if ctx.opset < 5:
-                # Old reshape takes new shape as attribute.
-                reshape = ctx.insert_new_node_on_input(node, "Reshape", kernel_name)
-                reshape.set_attr("shape", new_kernel_shape)
-                reshape.skip_conversion = True
+            if node.inputs[1].is_const():
+                input_node = node.inputs[1]
+                val = input_node.get_tensor_value(as_list=False)
+                val = np.reshape(val, new_kernel_shape)
+                input_node.set_tensor_value(val)
             else:
-                # New reshape takes new shape as input[1].
-                shape_name = utils.make_name(node.name)
-                ctx.make_const(shape_name, np.array(new_kernel_shape, dtype=np.int64))
-
-                reshape = ctx.make_node("Reshape", [kernel_name, shape_name])
-                ctx.replace_input(node, kernel_name, reshape.output[0], 1)
-
-                reshape.skip_conversion = True
-
-            ctx.set_shape(reshape.output[0], new_kernel_shape)
+                kernel_name = node.input[1]
+                if ctx.opset < 5:
+                    # Old reshape takes new shape as attribute.
+                    reshape = ctx.insert_new_node_on_input(node, "Reshape", kernel_name)
+                    reshape.set_attr("shape", new_kernel_shape)
+                    reshape.skip_conversion = True
+                else:
+                    # New reshape takes new shape as input[1].
+                    shape_name = utils.make_name(node.name)
+                    ctx.make_const(shape_name, np.array(new_kernel_shape, dtype=np.int64))
+
+                    reshape = ctx.make_node("Reshape", [kernel_name, shape_name])
+                    ctx.replace_input(node, kernel_name, reshape.output[0], 1)
+
+                    reshape.skip_conversion = True
+                ctx.set_shape(reshape.output[0], new_kernel_shape)
 
         # Get kernel (may have be changed to a reshape above).
         kernel_node = node.inputs[1]