Merge pull request #501 from zhijxu-MS/tmp_branch_for_PR2

enchace fold_const_optimizer

Author: nbcsm

tests/test_backend.py

@@ -1417,6 +1417,7 @@ def test_cancel_transpose(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    @check_onnxruntime_min_version("0.5.0", "topk-10's shape inference function has a bug")
     @check_opset_min_version(6, "cast")
     def test_topk1(self):
         x_val = np.arange(3 * 2 * 3).astype("float32")

tests/test_optimizers.py

@@ -484,6 +484,44 @@ 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)
+
+    def test_const_fold_cast_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("Cast", ["const"], ["value1"], to=TensorProto.INT64)
+        node3 = helper.make_node("Add", ["value1", "X"], ["res"])
+
+        graph = helper.make_graph(
+            [node1, node2, node3],
+            "test_const_fold_cast_with_const",
+            [helper.make_tensor_value_info("X", TensorProto.INT64, shape)],
+            [helper.make_tensor_value_info("res", TensorProto.INT64, shape)],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_and_compare(["res"], {"X": np.random.randn(*shape).astype(np.int64)}, model_proto,
+                             "Cast", 0)
     # Const Fold Optimizer Tests End
 
 

tf2onnx/optimizer/const_fold_optimizer.py

@@ -92,6 +92,14 @@ def _replace_node_with_const(node, graph, vals):
             graph.replace_all_inputs(graph.get_nodes(), old_input, const_node.output[0])
         graph.remove_node(node.name)
 
+    @staticmethod
+    @_register_func("Cast")
+    def _fold_cast(node, graph):
+        const_val = node.inputs[0].get_tensor_value(as_list=False)
+        np_dtype = utils.ONNX_TO_NUMPY_DTYPE[node.get_attr("to").i]
+        const_val_after_cast = const_val.astype(np_dtype)
+        return [const_val_after_cast]
+
     @staticmethod
     @_register_func("Transpose")
     def _fold_transpose(node, graph) -> list:
@@ -100,3 +108,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]