Add opset_version in export_onnx. (#21541)

Author: james77777778

keras/src/export/onnx.py

@@ -11,7 +11,14 @@
 from keras.src.utils import io_utils
 
 
-def export_onnx(model, filepath, verbose=None, input_signature=None, **kwargs):
+def export_onnx(
+    model,
+    filepath,
+    verbose=None,
+    input_signature=None,
+    opset_version=None,
+    **kwargs,
+):
     """Export the model as a ONNX artifact for inference.
 
     This method lets you export a model to a lightweight ONNX artifact
@@ -31,6 +38,9 @@ def export_onnx(model, filepath, verbose=None, input_signature=None, **kwargs):
             inputs. Can be a structure of `keras.InputSpec`, `tf.TensorSpec`,
             `backend.KerasTensor`, or backend tensor. If not provided, it will
             be automatically computed. Defaults to `None`.
+        opset_version: Optional. An integer value that specifies the ONNX opset
+            version. If not provided, the default version for the backend will
+            be used. Defaults to `None`.
         **kwargs: Additional keyword arguments.
 
     **Note:** This feature is currently supported only with TensorFlow, JAX and
@@ -82,7 +92,10 @@ def export_onnx(model, filepath, verbose=None, input_signature=None, **kwargs):
         # Use `tf2onnx` to convert the `decorated_fn` to the ONNX format.
         patch_tf2onnx()  # TODO: Remove this once `tf2onnx` supports numpy 2.
         tf2onnx.convert.from_function(
-            decorated_fn, input_signature, output_path=filepath
+            decorated_fn,
+            input_signature,
+            opset=opset_version,
+            output_path=filepath,
         )
 
     elif backend.backend() == "torch":
@@ -126,7 +139,11 @@ def export_onnx(model, filepath, verbose=None, input_signature=None, **kwargs):
             try:
                 # Try the TorchDynamo-based ONNX exporter first.
                 onnx_program = torch.onnx.export(
-                    model, sample_inputs, verbose=actual_verbose, dynamo=True
+                    model,
+                    sample_inputs,
+                    verbose=actual_verbose,
+                    opset_version=opset_version,
+                    dynamo=True,
                 )
                 if hasattr(onnx_program, "optimize"):
                     onnx_program.optimize()  # Only supported by torch>=2.6.0.
@@ -139,7 +156,11 @@ def export_onnx(model, filepath, verbose=None, input_signature=None, **kwargs):
 
                 # Fall back to the TorchScript-based ONNX exporter.
                 torch.onnx.export(
-                    model, sample_inputs, filepath, verbose=actual_verbose
+                    model,
+                    sample_inputs,
+                    filepath,
+                    verbose=actual_verbose,
+                    opset_version=opset_version,
                 )
     else:
         raise NotImplementedError(

keras/src/export/onnx_test.py

@@ -245,3 +245,27 @@ def build(self, y_shape, x_shape):
             )
         }
         ort_session.run(None, ort_inputs)
+
+    @parameterized.named_parameters(named_product(opset_version=[None, 18]))
+    def test_export_with_opset_version(self, opset_version):
+        import onnx as onnx_lib
+
+        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
+        model = get_model("sequential")
+        batch_size = 3 if backend.backend() != "torch" else 1
+        ref_input = np.random.normal(size=(batch_size, 10))
+        ref_input = ref_input.astype("float32")
+        ref_output = model(ref_input)
+
+        onnx.export_onnx(
+            model, temp_filepath, opset_version=opset_version, verbose=True
+        )
+        ort_session = onnxruntime.InferenceSession(temp_filepath)
+        ort_inputs = {
+            k.name: v for k, v in zip(ort_session.get_inputs(), [ref_input])
+        }
+        self.assertAllClose(ref_output, ort_session.run(None, ort_inputs)[0])
+
+        if opset_version is not None:
+            onnx_model = onnx_lib.load(temp_filepath)
+            self.assertEqual(onnx_model.opset_import[0].version, opset_version)

keras/src/models/model.py

@@ -548,17 +548,20 @@ def export(
                 `tf.TensorSpec`, `backend.KerasTensor`, or backend tensor. If
                 not provided, it will be automatically computed. Defaults to
                 `None`.
-            **kwargs: Additional keyword arguments:
-                - Specific to the JAX backend and `format="tf_saved_model"`:
-                    - `is_static`: Optional `bool`. Indicates whether `fn` is
-                        static. Set to `False` if `fn` involves state updates
-                        (e.g., RNG seeds and counters).
-                    - `jax2tf_kwargs`: Optional `dict`. Arguments for
-                        `jax2tf.convert`. See the documentation for
-                        [`jax2tf.convert`](
-                            https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md).
-                        If `native_serialization` and `polymorphic_shapes` are
-                        not provided, they will be automatically computed.
+            **kwargs: Additional keyword arguments.
+                - `is_static`: Optional `bool`. Specific to the JAX backend and
+                    `format="tf_saved_model"`. Indicates whether `fn` is static.
+                    Set to `False` if `fn` involves state updates (e.g., RNG
+                    seeds and counters).
+                - `jax2tf_kwargs`: Optional `dict`. Specific to the JAX backend
+                    and `format="tf_saved_model"`. Arguments for
+                    `jax2tf.convert`. See the documentation for
+                    [`jax2tf.convert`](
+                        https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md).
+                    If `native_serialization` and `polymorphic_shapes` are not
+                    provided, they will be automatically computed.
+                - `opset_version`: Optional `int`. Specific to `format="onnx"`.
+                    An integer value that specifies the ONNX opset version.
 
         **Note:** This feature is currently supported only with TensorFlow, JAX
         and Torch backends.