Implement const dequantize pushing for per-axis dequantization (#1321)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_optimizers.py

@@ -1512,6 +1512,86 @@ def test_const_dequantize_reshape(self):
         model_proto = self.make_model(graph, producer_name="onnx-tests")
         self.run_and_compare(["Z"], {}, model_proto, "Reshape", 0)
 
+    @check_opset_min_version(13, "DequantizeLinear")
+    def test_const_dequantize_reshape_per_channel(self):
+        inputval = numpy_helper.from_array(np.random.randint(0, 100, (2, 3, 4, 5), np.uint8), name='X')
+        scale = numpy_helper.from_array(np.array([0.75, 1., 0.2], dtype=np.float32), name='scale')
+        zero_point = numpy_helper.from_array(np.array([3, 4, 50], dtype=np.uint8), name='zero_point')
+        shape = numpy_helper.from_array(np.array([1, 1, 2, 3, 20], dtype=np.int64), name='shape')
+        node1 = helper.make_node("DequantizeLinear", ["X", "scale", "zero_point"], ["Y"], name="dequantize", axis=-3)
+        node2 = helper.make_node("Reshape", ["Y", "shape"], ["Z"], name="reshape")
+
+        graph = helper.make_graph(
+            [node1, node2],
+            "const-dequantize-test",
+            [],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 1, 2, 3, 20))],
+            [inputval, scale, zero_point, shape]
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_and_compare(["Z"], {}, model_proto, "Reshape", 0)
+
+    @check_opset_min_version(13, "DequantizeLinear")
+    def test_const_dequantize_reshape_per_channel_skipped(self):
+        inputval = numpy_helper.from_array(np.random.randint(0, 100, (2, 3, 4, 5), np.uint8), name='X')
+        scale = numpy_helper.from_array(np.array([0.75, 1., 0.2, 0.3], dtype=np.float32), name='scale')
+        zero_point = numpy_helper.from_array(np.array([3, 4, 50, 2], dtype=np.uint8), name='zero_point')
+        shape = numpy_helper.from_array(np.array([1, 6, 2, 2, 5], dtype=np.int64), name='shape')
+        node1 = helper.make_node("DequantizeLinear", ["X", "scale", "zero_point"], ["Y"], name="dequantize", axis=2)
+        node2 = helper.make_node("Reshape", ["Y", "shape"], ["Z"], name="reshape")
+
+        graph = helper.make_graph(
+            [node1, node2],
+            "const-dequantize-test",
+            [],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 6, 2, 2, 5))],
+            [inputval, scale, zero_point, shape]
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        # No optimization can be done here since the channel axis has changed size
+        self.run_and_compare(["Z"], {}, model_proto, "Reshape", 1)
+
+    @check_opset_min_version(13, "DequantizeLinear")
+    def test_const_dequantize_transpose_per_channel(self):
+        inputval = numpy_helper.from_array(np.random.randint(0, 100, (2, 3, 4, 5), np.uint8), name='X')
+        scale = numpy_helper.from_array(np.array([0.75, 1., 0.2], dtype=np.float32), name='scale')
+        zero_point = numpy_helper.from_array(np.array([3, 4, 50], dtype=np.uint8), name='zero_point')
+        node1 = helper.make_node("DequantizeLinear", ["X", "scale", "zero_point"], ["Y"], name="dequantize", axis=1)
+        node2 = helper.make_node("Transpose", ["Y"], ["Z"], name="transpose", perm=[0, 2, 3, 1])
+
+        graph = helper.make_graph(
+            [node1, node2],
+            "const-dequantize-test",
+            [],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 4, 5, 3))],
+            [inputval, scale, zero_point]
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_and_compare(["Z"], {}, model_proto, "Transpose", 0)
+
+    @check_opset_min_version(13, "DequantizeLinear")
+    def test_const_dequantize_unsqueeze_per_channel(self):
+        inputval = numpy_helper.from_array(np.random.randint(0, 100, (2, 3, 4, 5), np.uint8), name='X')
+        scale = numpy_helper.from_array(np.array([0.75, 1., 0.2], dtype=np.float32), name='scale')
+        zero_point = numpy_helper.from_array(np.array([3, 4, 50], dtype=np.uint8), name='zero_point')
+        axes = numpy_helper.from_array(np.array([-1, 0, -8, 3, 5], dtype=np.int64), name='axes')
+        node1 = helper.make_node("DequantizeLinear", ["X", "scale", "zero_point"], ["Y"], name="dequantize", axis=1)
+        node2 = helper.make_node("Unsqueeze", ["Y", "axes"], ["Z"], name="unsqueeze")
+
+        graph = helper.make_graph(
+            [node1, node2],
+            "const-dequantize-test",
+            [],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 1, 2, 1, 3, 1, 4, 5, 1))],
+            [inputval, scale, zero_point, axes]
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_and_compare(["Z"], {}, model_proto, "Transpose", 0)
+
     # Const Dequantize Optimizer Tests End
 
     def test_transpose_back_to_back_non_const(self):

tf2onnx/optimizer/const_dequantize_optimizer.py

@@ -45,6 +45,14 @@ def _fold_node(self, node, graph):
             return False
         if not self._all_inputs_are_const(dequant_node.inputs):
             return False
+        if len(dequant_node.inputs[1].get_tensor_value(as_list=False).flatten()) != 1:
+            # If using per-channel quantization, we must compute the new axis
+            old_axis = dequant_node.get_attr_value("axis")
+            input_shape = dequant_node.inputs[0].get_tensor_value(as_list=False).shape
+            new_axis = self.compute_new_axis(node, graph, old_axis, input_shape)
+            if new_axis is None:
+                return False
+            dequant_node.set_attr("axis", new_axis)
         graph.replace_input(node, node.input[0], dequant_node.input[0], 0)
         const_outputs = ConstFoldOptimizer.compute_const_folding(node, graph)
         graph.replace_all_inputs(node.output[0], dequant_node.output[0])
@@ -65,3 +73,40 @@ def _is_graph_output(node, graph):
         node_out_set = set(node.output)
         graph_out_set = set(graph.outputs)
         return node_out_set.intersection(graph_out_set)
+
+    @staticmethod
+    def compute_new_axis(node, graph, old_axis, input_shape):
+        if old_axis < 0:
+            old_axis += len(input_shape)
+        if node.type == "Transpose":
+            perm = node.get_attr_value("perm")
+            if perm is None:
+                return None
+            return perm.index(old_axis)
+        if node.type == "Reshape":
+            prod = 1
+            for d in input_shape[:old_axis+1]:
+                prod *= d
+            new_shape = node.inputs[1].get_tensor_value(as_list=True)
+            new_prod = 1
+            for i, d in enumerate(new_shape):
+                new_prod *= d
+                if new_prod == prod:
+                    if new_shape[i] == input_shape[old_axis]:
+                        return i
+                    return None
+            return None
+        if node.type == "Unsqueeze":
+            if graph.opset >= 13:
+                axes = node.inputs[1].get_tensor_value(as_list=True)
+            else:
+                axes = node.get_attr_value("axes")
+            new_rank = len(input_shape) + len(axes)
+            axes = [axis if axis >= 0 else axis + new_rank for axis in axes]
+            for i in range(new_rank):
+                if i not in axes:
+                    if old_axis == 0:
+                        return i
+                    old_axis -= 1
+            return None
+        return None