Fix code to enable opset 18 and latest onnxruntime. (#2143)

1. Update Split attributes for opset 18.
2. Update all Reduce* ops for opset 18.
3. Update transpose_optimizer for opset 18.
4. Enable opset 18 in CI tests.
---------

Signed-off-by: Jay Zhang <[email protected]>

Author: fatcat-z

ci_build/azure_pipelines/templates/job_generator.yml

@@ -4,9 +4,9 @@ parameters:
   platforms: ['linux']
   python_versions: ['3.7']
   tf_versions: ['']
-  onnx_versions: ['']
-  onnx_opsets: ['17', '16', '15', '14', '13']
-  onnx_backends: {onnxruntime: ['1.13.1']}
+  onnx_versions: ['1.13.1']
+  onnx_opsets: ['18', '17', '16', '15', '14', '13']
+  onnx_backends: {onnxruntime: ['1.14.1']}
   job: {}
   run_setup: 'True'
   report_coverage: 'False'

ci_build/azure_pipelines/templates/setup.yml

@@ -4,7 +4,8 @@ steps:
 - bash: |
     set -ex
     pip install pytest pytest-cov pytest-runner coverage graphviz requests pyyaml pillow pandas parameterized sympy coloredlogs flatbuffers timeout-decorator
-    pip install $(CI_PIP_TF_NAME) $(CI_PIP_ONNX_NAME)
+    pip install $(CI_PIP_TF_NAME)
+    pip install $(CI_PIP_ONNX_NAME)
 
     # TF < 2.7 reuires numpy <= 1.19, but onnxruntime >= 1.11 requires numpy >= 1.21
     if [[ $CI_TF_VERSION < 2.7 ]] && [[ $CI_ONNX_BACKEND == "onnxruntime" ]] ;

ci_build/azure_pipelines/templates/unit_test.yml

@@ -1,7 +1,7 @@
 # Run unit test
 
 parameters:
-    onnx_opsets: ['17', '16', '15', '14', '13']
+    onnx_opsets: ['18', '17', '16', '15', '14', '13']
     skip_tflite_tests: 'True'
     skip_tfjs_tests: 'True'
     skip_tf_tests: 'False'

ci_build/azure_pipelines/trimmed_keras2onnx_application_tests.yml

@@ -51,7 +51,7 @@ jobs:
         INSTALL_KERAS:
         UNINSTALL_KERAS:
         INSTALL_TENSORFLOW: pip install tensorflow==2.11.0
-        INSTALL_ORT: pip install onnxruntime==1.13.1
+        INSTALL_ORT: pip install onnxruntime==1.14.1
         INSTALL_KERAS_RESNET: pip install keras-resnet
         INSTALL_TRANSFORMERS: pip install transformers==4.2.0
         INSTALL_NUMPY:

ci_build/azure_pipelines/trimmed_keras2onnx_unit_test.yml

@@ -61,14 +61,14 @@ jobs:
         python.version: '3.8'
         ONNX_PATH: onnx==1.12.0
         TENSORFLOW_PATH: tensorflow==2.9.0
-        INSTALL_ORT: pip install onnxruntime==1.12.0
+        INSTALL_ORT: pip install onnxruntime==1.13.1
         INSTALL_NUMPY:
 
       Python310-tf-latest:
         python.version: '3.10'
         ONNX_PATH: onnx==1.13.1
         TENSORFLOW_PATH: tensorflow==2.11.0
-        INSTALL_ORT: pip install onnxruntime==1.13.1
+        INSTALL_ORT: pip install onnxruntime==1.14.1
         INSTALL_NUMPY:
 
       ############ Pure Keras Unit Tests ############

tests/keras2onnx_unit_tests/test_utils.py

@@ -20,7 +20,7 @@
 # https://github.com/onnx/tensorflow-onnx/issues/2132
 ORT_OPSET_VERSION = {
     "1.6.0": 13, "1.7.0": 13, "1.8.0": 14, "1.9.0": 15, "1.10.0": 15, "1.11.0": 16,
-    "1.12.0": 17, "1.13.0": 17, "1.14.0": 17
+    "1.12.0": 17, "1.13.0": 17, "1.14.0": 18
 }
 
 working_path = os.path.abspath(os.path.dirname(__file__))

tests/test_onnx_shape_inference.py

@@ -87,7 +87,7 @@ def test_split(self):
         shapes = [[5, 6, 7]]
         dtypes = [TensorProto.FLOAT]
         graph = self._create_empty_graph(inputs, shapes, dtypes)
-        node = graph.make_node("Split", [INPUT1], output_count=2, attr={"axis": 1})
+        node = graph.make_node("Split", [INPUT1], output_count=2, attr={"axis": 1, "num_outputs": 2})
         graph.add_graph_output(node.output[0])
         graph.add_graph_output(node.output[1])
         self._run_test_case(graph, self._generate_random_inputs(inputs, shapes, dtypes))

tests/test_optimizers.py

@@ -127,7 +127,10 @@ def test_transpose_with_split(self, input_shape, perm, inner_perm):
         output_shape = [output_before_trans[i] for i in perm]
         for axis in range(len(input_shape)):
             node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=inner_perm, name="trans1")
-            node2 = helper.make_node("Split", ["Y"], ["Z"], axis=axis, name="split")
+            if self.config.opset < 18:
+                node2 = helper.make_node("Split", ["Y"], ["Z"], axis=axis, name="split")
+            else:
+                node2 = helper.make_node("Split", ["Y"], ["Z"], axis=axis, name="split", num_outputs=1)
             node3 = helper.make_node("Transpose", ["Z"], ["res"], perm=perm, name="trans2")
 
             graph = helper.make_graph(
@@ -148,7 +151,10 @@ def test_transpose_with_split(self, input_shape, perm, inner_perm):
     @check_opset_max_version(12, "split attribute changed to input since opset 13")
     def test_transpose_with_split_dynamic_shape(self, input_shape, specific_input, output_shape, perm):
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm, name="trans")
-        node2 = helper.make_node("Split", ["Y"], ["Z"], axis=1, split=[1], name="split")
+        if self.config.opset < 18:
+            node2 = helper.make_node("Split", ["Y"], ["Z"], axis=1, split=[1], name="split")
+        else:
+            node2 = helper.make_node("Split", ["Y"], ["Z"], axis=1, split=[1], name="split", num_outputs=1)
         node3 = helper.make_node("Squeeze", ["Z"], ["B"], name="squeeze")
 
         graph = helper.make_graph(
@@ -1305,6 +1311,7 @@ def test_transpose_reciprocal(self, shape, perm_input, perm_output):
         ((1, 3, 4, 5), (1, 3, 1, 1), [0, 2, 3, 1], [0, 3, 1, 2]),
         ((1, 3, 4, 5, 6), (1, 3, 1, 1, 1), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])
+    @check_opset_max_version(17, "ReduceMean from opset <= 17 has axes as an attribute")
     def test_transpose_reducemean(self, input_shape, output_shape, perm_input, perm_output):
         node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
         node1 = helper.make_node("ReduceMean", ["Y"], ["Z"], axes=list(range(1, len(input_shape) - 1)),
@@ -1322,6 +1329,29 @@ def test_transpose_reducemean(self, input_shape, output_shape, perm_input, perm_
         self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    @parameterized.expand([
+        ((3, 4, 5), (3, 4, 1), [1], [0, 2, 1], [0, 2, 1]),
+        ((1, 3, 4, 5), (1, 3, 1, 1), [1, 2], [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5, 6), (1, 3, 1, 1, 1), [1, 2, 3], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+    ])
+    @check_opset_min_version(18, "ReduceMean from opset > 17 has axes as an input")
+    def test_transpose_reducemean_18(self, input_shape, output_shape, axes, perm_input, perm_output):
+        node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
+        node1 = helper.make_node("ReduceMean", ["Y", "axes"], ["Z"], keepdims=1, name="reducemean")
+        node2 = helper.make_node("Transpose", ["Z"], ["res"], perm=perm_output, name="trans_2")
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "transpose-reducemean-test-18",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, output_shape)],
+            [helper.make_tensor("axes", data_type=TensorProto.INT64,
+                                dims=[len(axes)], vals=axes)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
     @parameterized.expand([
         ((3, 4, 5), (3, 4, 1), [1], [0, 2, 1], [0, 2, 1]),
         ((1, 3, 4, 5), (1, 3, 4, 1), [2], [0, 2, 3, 1], [0, 3, 1, 2]),
@@ -1357,6 +1387,7 @@ def test_transpose_reducesum(self, input_shape, output_shape, axes, perm_input,
         ((1, 3, 4, 5, 6), (1, 3), [1, 2, 3], [0, 2, 3, 4, 1], [0, 1]),
         ((1, 3, 4, 5, 6), (), [0, 1, 2, 3, 4], [0, 2, 3, 4, 1], []),
     ])
+    @check_opset_max_version(17, "ReduceMax from opset <= 17 has axes as an attribute")
     def test_transpose_reducemax(self, input_shape, output_shape, axes, perm_input, perm_output):
         node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
         node1 = helper.make_node("ReduceMax", ["Y"], ["Z"], axes=axes,
@@ -1377,6 +1408,30 @@ def test_transpose_reducemax(self, input_shape, output_shape, axes, perm_input,
         self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    @parameterized.expand([
+        ((3, 4, 5), (3, 4, 1), [1], [0, 2, 1], [0, 2, 1]),
+        ((1, 3, 4, 5), (1, 3, 1, 1), [1, 2], [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5, 6), (1, 3, 1, 1, 1), [1, 2, 3], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+    ])
+    @check_opset_min_version(18, "ReduceMax from opset > 17 has axes as an input")
+    def test_transpose_reducemax_18(self, input_shape, output_shape, axes, perm_input, perm_output):
+        node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
+        node1 = helper.make_node("ReduceMax", ["Y", "axes"], ["Z"], keepdims=1, name="reducemax")
+        node2 = helper.make_node("Transpose", ["Z"], ["res"], perm=perm_output, name="trans_2")
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "transpose-reducemax-test-18",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, output_shape)],
+            [helper.make_tensor("axes", data_type=TensorProto.INT64,
+                                dims=[len(axes)], vals=axes)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
     def test_transpose_argmax(self):
         input_shape = [1, 2, 3, 4]
         node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans_1")
@@ -1858,6 +1913,7 @@ def test_duplicated_duplicated_input(self):
         self.run_merge_duplicated_nodes_compare(["OUT"], {"X": np.random.randn(5, 5).astype(np.float32)}, model_proto,
                                                 op_type="Add", remaining_op_num=2)
 
+    @check_opset_max_version(17, "ReduceMin from opset <= 17 has axes as an attribute")
     def test_duplicated_duplicated_attributes(self):
         # same attr or not
         node0 = helper.make_node('ReduceMin', inputs=["X"], outputs=["value0"], axes=[0], keepdims=0)
@@ -1877,6 +1933,31 @@ def test_duplicated_duplicated_attributes(self):
         self.run_merge_duplicated_nodes_compare(["OUT"], {"X": np.random.randn(5, 5).astype(np.float32)}, model_proto,
                                                 op_type="ReduceMin", remaining_op_num=2)
 
+    @check_opset_min_version(18, "ReduceMin from opset > 17 has axes as an input")
+    def test_duplicated_duplicated_attributes_18(self):
+        # same attr or not
+        node0 = helper.make_node('ReduceMin', inputs=["X", "axes0"], outputs=["value0"], keepdims=0)
+        node1 = helper.make_node('ReduceMin', inputs=["X", "axes1"], outputs=["value1"], keepdims=0)
+        node2 = helper.make_node('ReduceMin', inputs=["X", "axes2"], outputs=["value2"], keepdims=0)
+        node3 = helper.make_node('Add', inputs=["value0", "value1"], outputs=["value3"])
+        node4 = helper.make_node("Mul", ["value2", "value3"], ["OUT"])
+
+        graph = helper.make_graph(
+            [node0, node1, node2, node3, node4],
+            "test_duplicated_duplicated_attributes_18",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (5, 5))],
+            [helper.make_tensor_value_info("OUT", TensorProto.FLOAT, (5,))],
+            [
+                helper.make_tensor("axes0", data_type=TensorProto.INT64, dims=[1], vals=[0]),
+                helper.make_tensor("axes1", data_type=TensorProto.INT64, dims=[1], vals=[0]),
+                helper.make_tensor("axes2", data_type=TensorProto.INT64, dims=[1], vals=[1]),
+            ]
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_merge_duplicated_nodes_compare(["OUT"], {"X": np.random.randn(5, 5).astype(np.float32)}, model_proto,
+                                                op_type="ReduceMin", remaining_op_num=2)
+
     def _check_initializer_num(self, graph_proto, num):
         return num == len(graph_proto.initializer)
 
@@ -2062,7 +2143,10 @@ def test_reshape_opt_with_mul(self):
         node55 = helper.make_node("Constant", [], ["six"], value=six_tensor)
 
         node6 = helper.make_node("Gather", ["S", "g_indices"], ["dims12"])
-        node7 = helper.make_node("ReduceProd", ["dims12"], ["dims12_prod"], axes=[0])
+        if self.config.opset >= 18:
+            node7 = helper.make_node("ReduceProd", ["dims12", "axes"], ["dims12_prod"])
+        else:
+            node7 = helper.make_node("ReduceProd", ["dims12"], ["dims12_prod"], axes=[0])
         if self.config.opset >= 10:
             node8 = helper.make_node("Slice", ["S", "starts", "ends", ""], ["dim0"])
         else:
@@ -2312,7 +2396,10 @@ def test_const_fold_split(self):
         const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
                                           vals=np.random.randn(2, 6, 1).flatten().astype(np.float32))
         node0 = helper.make_node("Constant", [], ["const"], value=const_tensor)
-        node1 = helper.make_node("Split", ["const"], ["out1", "out2", "out3"], axis=1)
+        if self.config.opset < 18:
+            node1 = helper.make_node("Split", ["const"], ["out1", "out2", "out3"], axis=1)
+        else:
+            node1 = helper.make_node("Split", ["const"], ["out1", "out2", "out3"], axis=1, num_outputs=3)
         node2 = helper.make_node("Sum", ["inp", "out1", "out2", "out3"], ["out4"])
 
         graph = helper.make_graph(
@@ -2331,7 +2418,10 @@ def test_const_fold_split_one(self):
         const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
                                           vals=np.random.randn(2, 6, 1).flatten().astype(np.float32))
         node0 = helper.make_node("Constant", [], ["const"], value=const_tensor)
-        node1 = helper.make_node("Split", ["const"], ["out1"], axis=1)
+        if self.config.opset < 18:
+            node1 = helper.make_node("Split", ["const"], ["out1"], axis=1)
+        else:
+            node1 = helper.make_node("Split", ["const"], ["out1"], axis=1, num_outputs=1)
         node2 = helper.make_node("Sum", ["inp", "out1"], ["out4"])
 
         graph = helper.make_graph(
@@ -2374,7 +2464,11 @@ def test_const_fold_split_const_splits(self):
         const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
                                           vals=np.random.randn(2, 6, 1).flatten().astype(np.float32))
         node0 = helper.make_node("Constant", [], ["const"], value=const_tensor)
-        node2 = helper.make_node("Split", ["const"], ["out1", "out2", "out3"], axis=1, split=[1, 3, 2])
+        if self.config.opset < 18:
+            node2 = helper.make_node("Split", ["const"], ["out1", "out2", "out3"], axis=1, split=[1, 3, 2])
+        else:
+            node2 = helper.make_node("Split", ["const"], ["out1", "out2", "out3"], axis=1, split=[1, 3, 2],
+                                     num_outputs=3)
         node3 = helper.make_node("Sum", ["inp", "out2"], ["out4"])
 
         graph = helper.make_graph(

tf2onnx/constants.py

@@ -55,5 +55,5 @@
 # Note: opset 7 and opset 8 came out with IR3 but we need IR4 because of PlaceholderWithDefault
 # Refer from https://github.com/onnx/onnx/blob/main/docs/Versioning.md#released-versions
 OPSET_TO_IR_VERSION = {
-    1: 3, 2: 3, 3: 3, 4: 3, 5: 3, 6: 3, 7: 4, 8: 4, 9: 4, 10: 5, 11: 6, 12: 7, 13: 7, 14: 7, 15: 8, 16: 8, 17: 8
+    1: 3, 2: 3, 3: 3, 4: 3, 5: 3, 6: 3, 7: 4, 8: 4, 9: 4, 10: 5, 11: 6, 12: 7, 13: 7, 14: 7, 15: 8, 16: 8, 17: 8, 18: 8
 }