example with keras

Author: xadupre

examples/example_keras_tf2onnx.py

@@ -0,0 +1,40 @@
+import numpy as np
+import onnx
+import tensorflow as tf
+import tf2onnx
+import tf2onnx.keras2onnx_api
+
+
+class LeNet(tf.keras.Model):
+    def __init__(self):
+        super(LeNet, self).__init__()
+        self.conv2d_1 = tf.keras.layers.Conv2D(filters=6,
+                                               kernel_size=(3, 3), activation='relu',
+                                               input_shape=(32, 32, 1))
+        self.average_pool = tf.keras.layers.AveragePooling2D((3, 3))
+        self.conv2d_2 = tf.keras.layers.Conv2D(filters=16,
+                                               kernel_size=(3, 3), activation='relu')
+        self.flatten = tf.keras.layers.Flatten()
+        self.fc_1 = tf.keras.layers.Dense(120, activation='relu')
+        self.fc_2 = tf.keras.layers.Dense(84, activation='relu')
+        self.out = tf.keras.layers.Dense(10, activation='softmax')
+
+    def call(self, inputs, **kwargs):
+        x = self.conv2d_1(inputs)
+        x = self.average_pool(x)
+        x = self.conv2d_2(x)
+        x = self.average_pool(x)
+        x = self.flatten(x)
+        x = self.fc_2(self.fc_1(x))
+        return self.out(x)
+
+
+# Define a simple model
+model = LeNet()
+data = np.random.rand(2 * 416 * 416 * 3).astype(np.float32).reshape(2, 416, 416, 3)
+expected = model(data)
+
+# Get ConcreteFunction
+# concrete_func = tf.function(model).get_concrete_function(tf.TensorSpec([None, None, None, None], tf.float32))
+oxml = tf2onnx.keras2onnx_api.convert_keras(model, input_signature=[tf.TensorSpec([None, None, None, None], tf.float32)])
+onnx.save(oxml, "model.onnx")

tests/keras2onnx_unit_tests/test_subclassing.py

@@ -6,6 +6,7 @@
 import tensorflow as tf
 from test_utils import convert_keras_for_test as convert_keras
 from mock_keras2onnx.proto import is_tensorflow_older_than
+import tf2onnx
 
 if (not mock_keras2onnx.proto.is_tf_keras) or (not mock_keras2onnx.proto.tfcompat.is_tf2):
     pytest.skip("Tensorflow 2.0 only tests.", allow_module_level=True)
@@ -17,7 +18,7 @@ def __init__(self):
         self.conv2d_1 = tf.keras.layers.Conv2D(filters=6,
                                                kernel_size=(3, 3), activation='relu',
                                                input_shape=(32, 32, 1))
-        self.average_pool = tf.keras.layers.AveragePooling2D()
+        self.average_pool = tf.keras.layers.AveragePooling2D((3, 3))
         self.conv2d_2 = tf.keras.layers.Conv2D(filters=16,
                                                kernel_size=(3, 3), activation='relu')
         self.flatten = tf.keras.layers.Flatten()
@@ -91,8 +92,9 @@ def test_lenet(runner):
     lenet = LeNet()
     data = np.random.rand(2 * 416 * 416 * 3).astype(np.float32).reshape(2, 416, 416, 3)
     expected = lenet(data)
-    lenet._set_inputs(data)
-    oxml = convert_keras(lenet)
+    if hasattr(lenet, "_set_inputs"):
+        lenet._set_inputs(data)
+    oxml = convert_keras(lenet, input_signature=[tf.TensorSpec([None, None, None, None], tf.float32)])
     assert runner('lenet', oxml, data, expected)
 
 
@@ -234,15 +236,28 @@ def call(self, inputs, **kwargs):
     swm = Model()
     const_in = [tf.Variable([2, 4, 6, 8, 10], dtype=tf.int32, name="input")]
     expected = swm(const_in)
-    if hasattr(swm, "_set_input"):
-        swm._set_inputs(const_in)
-    else:
-        swm.inputs_spec = const_in
-    if hasattr(swm, "_set_output"):
-        swm._set_output(expected)
-    else:
-        swm.outputs_spec = expected
-    oxml = convert_keras(swm)
+
+    """
+    for op in concrete_func.graph.get_operations():
+        print("--", op.name)
+        print(op)
+
+    print("***", concrete_func.inputs)
+    print("***", concrete_func.outputs)
+    """
+    run_model = tf.function(swm)
+    concrete_func = run_model.get_concrete_function(tf.TensorSpec([None], tf.int32))
+    model_proto, external_tensor_storage = tf2onnx.convert._convert_common(
+        concrete_func.graph.as_graph_def(),
+        input_names=[i.name for i in concrete_func.inputs],
+        output_names=[i.name for i in concrete_func.outputs],
+        large_model=False,
+        output_path="where_test.onnx",
+    )
+    assert model_proto
+    assert not external_tensor_storage
+
+    oxml = convert_keras(swm, input_signature=[tf.TensorSpec([None], tf.int32)])
     assert runner('where_test', oxml, const_in, expected)
 
 

tf2onnx/convert.py

@@ -478,6 +478,7 @@ def from_keras3(model, input_signature=None, opset=None, custom_ops=None, custom
     reverse_lookup = {v: k for k, v in tensors_to_rename.items()}
 
     valid_names = []
+    model_output = None
     if hasattr(model, "outputs"):
         model_output = model.outputs
     else:
@@ -486,7 +487,7 @@ def from_keras3(model, input_signature=None, opset=None, custom_ops=None, custom
         elif model_input and len(model_input) == 1:
             # Let's try something to make unit test work. This should be replaced.
             model_output = [tf.Variable(model_input[0], name="output")]
-        else:
+        elif not output_names:
             raise RuntimeError(
                 "You should set attribute 'outputs_spec' with your outputs "
                 "so that the expected can use that information."
@@ -498,23 +499,20 @@ def _get_name(t, i):
         except AttributeError:
             return f"output:{i}"
 
-    for out in [_get_name(t, i) for i, t in enumerate(model_output)]:
-        if out in reverse_lookup:
-            valid_names.append(reverse_lookup[out])
-        else:
-            print(f"Warning: Output name '{out}' not found in reverse_lookup.")
-            # Fallback: verwende TensorFlow-Ausgangsnamen direkt
-            valid_names = [
-                _get_name(t, i)
-                for i, t in enumerate(concrete_func.outputs)
-                if t.dtype != tf.dtypes.resource
-            ]
-            break
-    output_names = valid_names
-
-
-    #if old_out_names is not None:
-        #model.output_names = old_out_names
+    if model_output:
+        for out in [_get_name(t, i) for i, t in enumerate(model_output)]:
+            if out in reverse_lookup:
+                valid_names.append(reverse_lookup[out])
+            else:
+                print(f"Warning: Output name '{out}' not found in reverse_lookup.")
+                # Fallback: verwende TensorFlow-Ausgangsnamen direkt
+                valid_names = [
+                    _get_name(t, i)
+                    for i, t in enumerate(concrete_func.outputs)
+                    if t.dtype != tf.dtypes.resource
+                ]
+                break
+        output_names = valid_names
 
     with tf.device("/cpu:0"):
         frozen_graph, initialized_tables = \

tf2onnx/keras2onnx_api.py

@@ -44,7 +44,8 @@ def get_maximum_opset_supported():
     return min(max(OPSET_TO_IR_VERSION.keys()), defs.onnx_opset_version())
 
 def convert_keras(model, name=None, doc_string='', target_opset=None, initial_types=None,
-                  channel_first_inputs=None, debug_mode=False, custom_op_conversions=None):
+                  channel_first_inputs=None, debug_mode=False, custom_op_conversions=None,
+                  input_signature=None):
     """
     :param model: keras model
     :param name: the converted onnx model internal name
@@ -54,16 +55,18 @@ def convert_keras(model, name=None, doc_string='', target_opset=None, initial_ty
     :param channel_first_inputs: A list of channel first input
     :param debug_mode: ignored
     :param custom_op_conversions: ignored
+    :param input_signature: takes precedence on initial_types if specified,
+        example: ``[tf.TensorSpec([None], tf.int32)]``
     :return an ONNX ModelProto
     """
     if target_opset is None:
         target_opset = get_maximum_opset_supported()
-    input_signature = _process_initial_types(initial_types, unknown_dim=None)
+    if input_signature is None:
+        input_signature = _process_initial_types(initial_types, unknown_dim=None)
     name = name or model.name
 
     model, _ = tf2onnx.convert.from_keras(model, input_signature, opset=target_opset,
                                           inputs_as_nchw=channel_first_inputs)
     model.graph.name = name
     model.graph.doc_string = doc_string
-
     return model

tf2onnx/onnx_opset/misc.py

@@ -42,3 +42,13 @@ class NukeNode:
     @classmethod
     def version_1(cls, ctx, node, **kwargs):
         ctx.remove_node(node.name)
+
+
+@tf_op("StatefulPartitionedCall")
+class StatefulPartitionedCall:
+    @classmethod
+    def version_1(cls, ctx, node, **kwargs):
+        raise NotImplementedError(
+            "This node appears if the graph has local function. It should be inlined first. "
+            "Inline did not work on that model. It seems Conv2D is no longer inlined."
+        )

tf2onnx/tf_loader.py

@@ -698,6 +698,19 @@ def tf_optimize_grappler(input_names, output_names, graph_def):
     from tensorflow.core.protobuf import meta_graph_pb2 as meta_graph_pb2, config_pb2, rewriter_config_pb2
     from tensorflow.python.grappler import tf_optimizer as tf_opt
 
+    rewriter_config = rewriter_config_pb2.RewriterConfig(
+        function_optimization=rewriter_config_pb2.RewriterConfig.ON
+    )
+
+    config = tf.compat.v1.ConfigProto()
+    config.graph_options.rewrite_options.CopyFrom(rewriter_config)
+
+    with tf.compat.v1.Session(config=config) as sess:
+        tf.import_graph_def(graph_def, name="")
+        optimized_graph_def = sess.graph.as_graph_def(add_shapes=True)
+
+    return optimized_graph_def
+
     config = config_pb2.ConfigProto()
     rewrite_options = config.graph_options.rewrite_options
     config.graph_options.infer_shapes = True