Fix the axis of inserted QDQ for ConvTranspose (#2134)

sdua-nv · fatcat-z · web-flow · commit ce29107603df · 2023-03-15T13:14:56.000Z
The quantization axis of QDQ nodes that are being inserted
before the kernel weights of all Conv nodes is currently 0.
This is incorrect; ConvTranspose requires axis=1.

Signed-off-by: Sirej Dua &lt;sdua@nvidia.com&gt;
Co-authored-by: Jay Zhang &lt;36183870+fatcat-z@users.noreply.github.com&gt;
diff --git a/tests/backend_test_base.py b/tests/backend_test_base.py
@@ -284,7 +284,10 @@ def run_tflite(self, tflite_path, feed_dict):
             # tflite sometimes converts from tf but produces an invalid model
             return None, None
 
-    def assert_shapes_correct(self, graph, allow_missing=False, run_checker=True):
+    def assert_shapes_correct(self, graph, allow_missing=False, run_checker=True, check_shape=True):
+        if not check_shape:
+            return None
+
         model_proto = graph.make_model("test")
 
         if run_checker and not any(graph.get_shape(out) is None for out in graph.outputs + graph.input_names):
@@ -402,8 +405,10 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
                     i = output_names_with_port.index(output_name)
                     actual[i] = np.transpose(actual[i], constants.NCHW_TO_NHWC)
 
-            self.assert_results_equal(expected, actual, rtol, atol, mtol, check_value, check_shape, check_dtype)
-            self.assert_shapes_correct(g, self.config.allow_missing_shapes, not self.config.skip_onnx_checker)
+            self.assert_results_equal(expected, actual, rtol, atol, mtol, check_value, check_shape,
+                                      check_dtype)
+            self.assert_shapes_correct(g, self.config.allow_missing_shapes, not self.config.skip_onnx_checker,
+                                       check_shape)
 
             if graph_validator:
                 self.assertTrue(graph_validator(g))
@@ -441,7 +446,8 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
                     onnx_tfl_res[i] = np.transpose(onnx_tfl_res[i], constants.NCHW_TO_NHWC)
 
             self.assert_results_equal(tfl_res, onnx_tfl_res, rtol, atol, mtol, check_value, check_shape, check_dtype)
-            self.assert_shapes_correct(g, self.config.allow_missing_shapes, not self.config.skip_onnx_checker)
+            self.assert_shapes_correct(g, self.config.allow_missing_shapes, not self.config.skip_onnx_checker,
+                                       check_shape)
 
             if graph_validator:
                 self.assertTrue(graph_validator(g))
@@ -475,7 +481,8 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
 
             self.assert_results_equal(tfjs_res, onnx_tfjs_res, rtol, atol, mtol, check_value, check_shape,
                                       check_dtype=False)
-            self.assert_shapes_correct(g, self.config.allow_missing_shapes, not self.config.skip_onnx_checker)
+            self.assert_shapes_correct(g, self.config.allow_missing_shapes, not self.config.skip_onnx_checker,
+                                       check_shape)
 
             if graph_validator:
                 self.assertTrue(graph_validator(g))
diff --git a/tests/common.py b/tests/common.py
@@ -51,6 +51,7 @@
     "check_op_count",
     "check_gru_count",
     "check_lstm_count",
+    "check_quantization_axis",
     "timeout",
 ]
 
@@ -471,6 +472,8 @@ def check_lstm_count(graph, expected_count):
 def check_gru_count(graph, expected_count):
     return check_op_count(graph, "GRU", expected_count)
 
+def check_quantization_axis(graph, op_type, expected_axis):
+    return np.all(np.array([n.get_attr_int("axis") for n in group_nodes_by_type(graph)[op_type]]) == expected_axis)
 
 _MAX_MS_OPSET_VERSION = 1
 
diff --git a/tests/test_backend.py b/tests/test_backend.py
@@ -607,6 +607,44 @@ def func(x, output_shape_placeholder):
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: output_shape},
                             rtol=1e-05, process_args=process_args)
 
+    @check_opset_min_version(10, "quantize_and_dequantize")
+    def test_conv2d_quantization_axis(self):
+        x_shape = [1, 1, 5, 5]
+        kernel_shape = _KERNEL3x3
+        strides = [1, 1, 1, 1]
+        x_val = make_xval(x_shape).transpose(NCHW_TO_NHWC)
+        kernel_val = make_xval(_KERNEL3x3)
+
+        def func(x):
+            f = tf.constant(kernel_val, name="kernel", dtype=tf.float32)
+            kernel_dq = quantize_and_dequantize(f, 0, np.prod(kernel_shape))
+            conv = tf.nn.conv2d(x, kernel_dq, strides=strides, padding="VALID")
+            return tf.identity(conv, name=_TFOUTPUT)
+        def graph_validator(g):
+            return check_quantization_axis(g, "DequantizeLinear", 0)
+
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, rtol=1e-05, graph_validator=graph_validator,
+                            check_shape=False)
+
+    @check_opset_min_version(10, "quantize_and_dequantize")
+    def test_conv2d_transpose_quantization_axis(self):
+        x_shape = [2, 6, 4, 3]
+        output_shape = [2, 13, 9, 2]
+        kernel_shape = [3, 3, 2, 3]
+        strides = [1, 2, 2, 1]
+        x_val = make_xval(x_shape)
+        kernel_val = make_xval(kernel_shape)
+        def func(x):
+            f = tf.constant(kernel_val, name="kernel", dtype=tf.float32)
+            kernel_dq = quantize_and_dequantize(f, 0, np.prod(kernel_shape))
+            conv = tf.nn.conv2d_transpose(x, kernel_dq, output_shape, strides=strides, padding="VALID")
+            return tf.identity(conv, name=_TFOUTPUT)
+        def graph_validator(g):
+            return check_quantization_axis(g, "DequantizeLinear", 1)
+
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, rtol=1e-05, graph_validator=graph_validator,
+                            check_shape=False)
+
     def test_depthwiseconv_0(self):
         x_shape = [1, 3, 4, 3]
         kernel_shape = [3, 3, 3, 3]
diff --git a/tf2onnx/onnx_opset/nn.py b/tf2onnx/onnx_opset/nn.py
@@ -58,7 +58,8 @@ def get_channels_last_permutation(spatial):
 
 
 def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
-                        input_indices=None, output_indices=None, spatial=2):
+                        input_indices=None, output_indices=None, spatial=2,
+                        quantization_axis=0):
     """Convert input and kernel from tensorflow to onnx. This may be required to
         insert transpose ops for input, kernel, and output unless they are constants
         and we can transpose the constant.
@@ -73,6 +74,7 @@ def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
         new_kernel_shape: Pass to reshape the kernel.
         input_indices: Indices that define the inputs.
         output_indices: Indices that define the outputs.
+        quantization_axis: Axis for the inserted QDQ nodes
     """
 
     if input_indices is None:
@@ -151,8 +153,8 @@ def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
                 weights_node.set_tensor_value(val)
                 need_transpose = False
                 # Change the quantization axis for Q and DQ node accordingly
-                kernel_node.set_attr("axis", 0) # DQ node
-                kernel_node.inputs[0].set_attr("axis", 0) # Q node
+                kernel_node.set_attr("axis", quantization_axis) # DQ node
+                kernel_node.inputs[0].set_attr("axis", quantization_axis) # Q node
             else:
                 val = kernel_node.get_tensor_value(as_list=False)
                 val = np.transpose(val, permutation)
@@ -607,7 +609,7 @@ def version_1(cls, ctx, node, **kwargs):
         ctx.replace_input(node, node.input[0], node.input[1], 0)
         ctx.replace_input(node, node.input[1], t, 1)
 
-        conv_convert_inputs(ctx, node, with_kernel=True, spatial=spatial)
+        conv_convert_inputs(ctx, node, with_kernel=True, spatial=spatial, quantization_axis=1)
 
     @classmethod
     def version_11(cls, ctx, node, **kwargs):
