enhance const_fold_optimizer: supports to fold const with Unsqueeze

Author: zhijxu-MS

tests/test_optimizers.py

@@ -484,6 +484,25 @@ def test_const_fold_node_is_output(self):
         model_proto = helper.make_model(graph, producer_name="onnx-tests")
         self.run_transpose_compare(["res"], {},
                                    model_proto, remaining_transpose_num=0)
+
+    def test_const_fold_unsqueeze_with_const(self):
+        shape = (6, 6)
+        const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
+                                          vals=np.random.randn(*shape).flatten().astype(np.float32))
+        node1 = helper.make_node("Constant", [], ["const"], value=const_tensor)
+        node2 = helper.make_node("Unsqueeze", ["const"], ["value1"], axes=[0, 2, 3])
+        node3 = helper.make_node("Add", ["value1", "X"], ["res"])
+
+        graph = helper.make_graph(
+            [node1, node2, node3],
+            "test_const_fold_unsqueeze_with_const",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1,))],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, (1, 6, 1, 1, 6))],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_and_compare(["res"], {"X": np.random.randn(1).astype(np.float32)}, model_proto,
+                             "Unsqueeze", 0)
     # Const Fold Optimizer Tests End
 
 

tf2onnx/optimizer/const_fold_optimizer.py

@@ -100,3 +100,27 @@ def _fold_transpose(node, graph) -> list:
         perm = perm_attr.ints if perm_attr else None
         const_val_after_trans = const_val.transpose(perm)
         return [const_val_after_trans]
+
+    @staticmethod
+    @_register_func("Unsqueeze")
+    def _fold_unsqueeze(node, graph):
+        """
+        numpy expand_dims only supports to unsqueeze one dim one time, so reshape is used to simplify the logic
+        """
+        const_val = node.inputs[0].get_tensor_value(as_list=False)
+        axes = list(node.get_attr("axes").ints)
+        utils.make_sure(all(axis >= 0 for axis in axes), "onnx spec says it only supports positive axis")
+        shape_in = const_val.shape
+        dims_out = len(shape_in) + len(axes)
+        # calculate the shape of output accroding to onnx Unsqueeze's spec
+        # https://github.com/onnx/onnx/blob/master/docs/Operators.md#Unsqueeze
+        shape_in = iter(shape_in)
+        shape_out = [None] * dims_out
+        for ind in axes:
+            shape_out[ind] = 1
+        for ind, val in enumerate(shape_out):
+            if val is None:
+                shape_out[ind] = next(shape_in)
+
+        const_val_after_unsqueeze = const_val.reshape(shape_out)
+        return [const_val_after_unsqueeze]