yobx.tensorflow.to_onnx#

yobx.tensorflow.to_onnx(model, args: ~typing.Tuple[~typing.Any], input_names: ~typing.Sequence[str] | None = None, dynamic_shapes: ~typing.Tuple[~typing.Dict[int, str]] | None = None, target_opset: int | ~typing.Dict[str, int] = 21, builder_cls: type | ~typing.Callable = <class 'yobx.xbuilder.graph_builder.GraphBuilder'>, verbose: int = 0, extra_converters: ~typing.Dict[str, ~typing.Callable] | None = None, large_model: bool = False, external_threshold: int = 1024) → ExportArtifact[source]#

Converts a TensorFlow/Keras model into ONNX.

The model is first traced with get_concrete_function() to obtain the actual TensorFlow computation graph. Each operation in that graph is then converted individually to an equivalent ONNX node by a registered converter.

Parameters:

  • model – a Keras model or layer (must be built / called at least once)

  • args – dummy inputs (numpy arrays or onnx.ValueInfoProto descriptors); used to determine dtypes and shapes. When a ValueInfoProto is provided no actual data is needed — the element type and shape are taken directly from the proto.

  • input_names – optional list of names for the ONNX input tensors

  • dynamic_shapes – optional per-input axis-to-dim-name mappings. When None, axis 0 is treated as a dynamic batch dimension for every input.

  • target_opset – opset to use; either an integer for the default domain (""), or a dictionary mapping domain names to opset versions. If it includes {'com.microsoft': 1}, the converted model may include optimized kernels specific to onnxruntime.

  • builder_cls – by default the graph builder is a yobx.xbuilder.GraphBuilder but any builder can be used as long it implements the apis Shape and type tracking and Building a graph from scratch

  • verbose – verbosity level (0 = silent)

  • extra_converters – optional mapping from TF op-type string to converter function with signature (g, sts, outputs, op, verbose=0); entries here take priority over the built-in op converters

  • large_model – if True the returned ExportArtifact has its container attribute set to an ExtendedModelContainer

  • external_threshold – if large_model is True, every tensor whose element count exceeds this threshold is stored as external data

Returns:

ExportArtifact wrapping the exported ONNX proto together with an ExportReport.

Example:

import numpy as np
import tensorflow as tf
from yobx.tensorflow import to_onnx

model = tf.keras.Sequential([tf.keras.layers.Dense(4, input_shape=(3,))])
X = np.random.randn(5, 3).astype(np.float32)

artifact = to_onnx(model, (X,))
proto = artifact.proto
artifact.save("model.onnx")
On this page

This Page