Merge branch 'support_export' into gsoc2025

Author: Mohamed-Ashraf273

keras/src/export/__init__.py

@@ -1,4 +1,5 @@
 from keras.src.export.onnx import export_onnx
+from keras.src.export.openvino import export_openvino
 from keras.src.export.saved_model import ExportArchive
 from keras.src.export.saved_model import export_saved_model
 from keras.src.export.tfsm_layer import TFSMLayer

keras/src/export/openvino.py

@@ -0,0 +1,107 @@
+from keras.src import backend
+from keras.src import tree
+from keras.src.export.export_utils import convert_spec_to_tensor
+from keras.src.export.export_utils import get_input_signature
+from keras.src.utils import io_utils
+
+
+def export_openvino(
+    model, filepath, verbose=None, input_signature=None, **kwargs
+):
+    """Export the model as an OpenVINO IR artifact for inference.
+
+    This method exports the model to the OpenVINO IR format,
+    which includes two files:
+    a `.xml` file containing the model structure and a `.bin` file
+    containing the weights.
+    The exported model contains only the forward pass
+    (i.e., the model's `call()` method), and can be deployed with the
+    OpenVINO Runtime for fast inference on CPU and other Intel hardware.
+
+    Args:
+        filepath: `str` or `pathlib.Path`. Path to the output `.xml` file.
+        The corresponding `.bin` file will be saved alongside it.
+        verbose: Optional `bool`. Whether to print a confirmation message
+        after export. If `None`, it uses the default verbosity configured
+        by the backend.
+        input_signature: Optional. Specifies the shape and dtype of the
+        model inputs. If not provided, it will be inferred.
+        **kwargs: Additional keyword arguments (currently unused).
+    """
+    assert filepath.endswith(".xml"), (
+        "The OpenVINO export requires the filepath to end with '.xml'. "
+        f"Got: {filepath}"
+    )
+
+    import openvino as ov
+    from openvino.runtime import opset14 as ov_opset
+
+    from keras.src.backend.openvino.core import OPENVINO_DTYPES
+    from keras.src.backend.openvino.core import OpenVINOKerasTensor
+
+    actual_verbose = verbose if verbose is not None else True
+
+    if input_signature is None:
+        input_signature = get_input_signature(model)
+    if isinstance(input_signature, list) and len(input_signature) == 1:
+        input_signature = input_signature[0]
+
+    sample_inputs = tree.map_structure(
+        lambda x: convert_spec_to_tensor(x, replace_none_number=1),
+        input_signature,
+    )
+
+    if backend.backend() == "openvino":
+        import inspect
+
+        def parameterize_inputs(inputs, prefix=""):
+            if isinstance(inputs, (list, tuple)):
+                return [
+                    parameterize_inputs(e, f"{prefix}{i}")
+                    for i, e in enumerate(inputs)
+                ]
+            elif isinstance(inputs, dict):
+                return {k: parameterize_inputs(v, k) for k, v in inputs.items()}
+            elif isinstance(inputs, OpenVINOKerasTensor):
+                ov_type = OPENVINO_DTYPES[str(inputs.dtype)]
+                ov_shape = list(inputs.shape)
+                param = ov_opset.parameter(shape=ov_shape, dtype=ov_type)
+                param.set_friendly_name(prefix)
+                return OpenVINOKerasTensor(param.output(0))
+            else:
+                raise TypeError(f"Unknown input type: {type(inputs)}")
+
+        params = parameterize_inputs(sample_inputs)
+        signature = inspect.signature(model.call)
+        if len(signature.parameters) > 1 and isinstance(params, (list, tuple)):
+            outputs = model(*params)
+        else:
+            outputs = model(params)
+        parameters = [p.output.get_node() for p in tree.flatten(params)]
+        results = [ov_opset.result(r.output) for r in tree.flatten(outputs)]
+        ov_model = ov.Model(results=results, parameters=parameters)
+        flat_specs = tree.flatten(input_signature)
+        for ov_input, spec in zip(ov_model.inputs, flat_specs):
+            # Respect the dynamic axes from the original input signature.
+            dynamic_shape_dims = [
+                -1 if dim is None else dim for dim in spec.shape
+            ]
+            dynamic_shape = ov.PartialShape(dynamic_shape_dims)
+            ov_input.get_node().set_partial_shape(dynamic_shape)
+
+    elif backend.backend() == "tensorflow":
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            model.export(temp_dir, format="tf_saved_model")
+            ov_model = ov.convert_model(temp_dir)
+    else:
+        raise NotImplementedError(
+            "`export_openvino` is only compatible with OpenVINO and "
+            "TensorFlow backends."
+        )
+
+    ov.serialize(ov_model, filepath)
+
+    if actual_verbose:
+        io_utils.print_msg(f"Saved OpenVINO IR at '{filepath}'.")

keras/src/export/openvino_test.py

@@ -0,0 +1,230 @@
+import os
+
+import numpy as np
+import pytest
+from absl.testing import parameterized
+
+from keras.src import backend
+from keras.src import layers
+from keras.src import models
+from keras.src import ops
+from keras.src import testing
+from keras.src import tree
+from keras.src.export import openvino
+from keras.src.saving import saving_lib
+from keras.src.testing.test_utils import named_product
+
+try:
+    import openvino as ov
+except ImportError:
+    ov = None
+
+
+class CustomModel(models.Model):
+    def __init__(self, layer_list):
+        super().__init__()
+        self.layer_list = layer_list
+
+    def call(self, input):
+        output = input
+        for layer in self.layer_list:
+            output = layer(output)
+        return output
+
+
+def get_model(type="sequential", input_shape=(10,), layer_list=None):
+    layer_list = layer_list or [
+        layers.Dense(10, activation="relu"),
+        layers.BatchNormalization(),
+        layers.Dense(1, activation="sigmoid"),
+    ]
+    if type == "sequential":
+        return models.Sequential(layer_list)
+    elif type == "functional":
+        input = output = tree.map_shape_structure(layers.Input, input_shape)
+        for layer in layer_list:
+            output = layer(output)
+        return models.Model(inputs=input, outputs=output)
+    elif type == "subclass":
+        return CustomModel(layer_list)
+    elif type == "lstm":
+        # https://github.com/keras-team/keras/issues/21390
+        inputs = layers.Input((4, 10))
+        x = layers.Bidirectional(
+            layers.LSTM(
+                10,
+                kernel_initializer="he_normal",
+                return_sequences=True,
+                kernel_regularizer=None,
+            ),
+            merge_mode="sum",
+        )(inputs)
+        outputs = layers.Bidirectional(
+            layers.LSTM(
+                10,
+                kernel_initializer="he_normal",
+                return_sequences=True,
+                kernel_regularizer=None,
+            ),
+            merge_mode="concat",
+        )(x)
+        return models.Model(inputs=inputs, outputs=outputs)
+
+
+@pytest.mark.skipif(ov is None, reason="OpenVINO is not installed")
+@pytest.mark.skipif(
+    backend.backend() not in ("tensorflow", "openvino"),
+    reason=(
+        "`export_openvino` only currently supports"
+        "the tensorflow and openvino backends."
+    ),
+)
+@pytest.mark.skipif(
+    testing.tensorflow_uses_gpu(), reason="Leads to core dumps on CI"
+)
+class ExportOpenVINOTest(testing.TestCase):
+    @parameterized.named_parameters(
+        named_product(
+            model_type=["sequential", "functional", "subclass", "lstm"]
+        )
+    )
+    def test_standard_model_export(self, model_type):
+        if model_type == "lstm":
+            self.skipTest(
+                "LSTM export not supported - unimplemented QR operation"
+            )
+
+        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model.xml")
+        model = get_model(model_type)
+        batch_size = 3
+        if model_type == "lstm":
+            ref_input = np.random.normal(size=(batch_size, 4, 10))
+        else:
+            ref_input = np.random.normal(size=(batch_size, 10))
+        ref_input = ref_input.astype("float32")
+        ref_output = model(ref_input)
+
+        openvino.export_openvino(model, temp_filepath)
+
+        # Load and run inference with OpenVINO
+        core = ov.Core()
+        ov_model = core.read_model(temp_filepath)
+        compiled_model = core.compile_model(ov_model, "CPU")
+
+        ov_output = compiled_model([ref_input])[compiled_model.output(0)]
+
+        self.assertAllClose(ref_output, ov_output)
+
+        larger_input = np.concatenate([ref_input, ref_input], axis=0)
+        compiled_model([larger_input])
+
+    @parameterized.named_parameters(
+        named_product(struct_type=["tuple", "array", "dict"])
+    )
+    def test_model_with_input_structure(self, struct_type):
+        class TupleModel(models.Model):
+            def call(self, inputs):
+                x, y = inputs
+                return ops.add(x, y)
+
+        class ArrayModel(models.Model):
+            def call(self, inputs):
+                x = inputs[0]
+                y = inputs[1]
+                return ops.add(x, y)
+
+        class DictModel(models.Model):
+            def call(self, inputs):
+                x = inputs["x"]
+                y = inputs["y"]
+                return ops.add(x, y)
+
+        batch_size = 3
+        ref_input = np.random.normal(size=(batch_size, 10)).astype("float32")
+        if struct_type == "tuple":
+            model = TupleModel()
+            ref_input = (ref_input, ref_input * 2)
+        elif struct_type == "array":
+            model = ArrayModel()
+            ref_input = [ref_input, ref_input * 2]
+        elif struct_type == "dict":
+            model = DictModel()
+            ref_input = {"x": ref_input, "y": ref_input * 2}
+
+        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model.xml")
+        ref_output = model(tree.map_structure(ops.convert_to_tensor, ref_input))
+
+        openvino.export_openvino(model, temp_filepath)
+
+        # Load and run inference with OpenVINO
+        core = ov.Core()
+        ov_model = core.read_model(temp_filepath)
+        compiled_model = core.compile_model(ov_model, "CPU")
+
+        if isinstance(ref_input, dict):
+            ov_inputs = [ref_input[key] for key in ref_input.keys()]
+        else:
+            ov_inputs = list(ref_input)
+
+        ov_output = compiled_model(ov_inputs)[compiled_model.output(0)]
+        self.assertAllClose(ref_output, ov_output)
+
+        # Test with keras.saving_lib
+        temp_filepath = os.path.join(
+            self.get_temp_dir(), "exported_model.keras"
+        )
+        saving_lib.save_model(model, temp_filepath)
+        revived_model = saving_lib.load_model(
+            temp_filepath,
+            {
+                "TupleModel": TupleModel,
+                "ArrayModel": ArrayModel,
+                "DictModel": DictModel,
+            },
+        )
+        self.assertAllClose(ref_output, revived_model(ref_input))
+        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model2.xml")
+        openvino.export_openvino(revived_model, temp_filepath)
+
+        bigger_ref_input = tree.map_structure(
+            lambda x: np.concatenate([x, x], axis=0), ref_input
+        )
+        if isinstance(bigger_ref_input, dict):
+            bigger_ov_inputs = [
+                bigger_ref_input[key] for key in bigger_ref_input.keys()
+            ]
+        else:
+            bigger_ov_inputs = list(bigger_ref_input)
+        compiled_model(bigger_ov_inputs)
+
+    def test_model_with_multiple_inputs(self):
+        class TwoInputsModel(models.Model):
+            def call(self, x, y):
+                return x + y
+
+            def build(self, y_shape, x_shape):
+                self.built = True
+
+        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model.xml")
+        model = TwoInputsModel()
+        batch_size = 3
+        ref_input_x = np.random.normal(size=(batch_size, 10)).astype("float32")
+        ref_input_y = np.random.normal(size=(batch_size, 10)).astype("float32")
+        ref_output = model(ref_input_x, ref_input_y)
+
+        openvino.export_openvino(model, temp_filepath)
+
+        # Load and run inference with OpenVINO
+        core = ov.Core()
+        ov_model = core.read_model(temp_filepath)
+        compiled_model = core.compile_model(ov_model, "CPU")
+
+        ov_output = compiled_model([ref_input_x, ref_input_y])[
+            compiled_model.output(0)
+        ]
+        self.assertAllClose(ref_output, ov_output)
+        larger_input_x = np.concatenate([ref_input_x, ref_input_x], axis=0)
+        larger_input_y = np.concatenate([ref_input_y, ref_input_y], axis=0)
+        compiled_model([larger_input_x, larger_input_y])
+        larger_input_y = np.concatenate([ref_input_y, ref_input_y], axis=0)
+        compiled_model([larger_input_x, larger_input_y])

keras/src/models/model.py

@@ -595,9 +595,10 @@ def export(
         ```
         """
         from keras.src.export import export_onnx
+        from keras.src.export import export_openvino
         from keras.src.export import export_saved_model
 
-        available_formats = ("tf_saved_model", "onnx")
+        available_formats = ("tf_saved_model", "onnx", "openvino")
         if format not in available_formats:
             raise ValueError(
                 f"Unrecognized format={format}. Supported formats are: "
@@ -620,6 +621,14 @@ def export(
                 input_signature=input_signature,
                 **kwargs,
             )
+        elif format == "openvino":
+            export_openvino(
+                self,
+                filepath,
+                verbose,
+                input_signature=input_signature,
+                **kwargs,
+            )
 
     @classmethod
     def from_config(cls, config, custom_objects=None):