yobx.container#

modules

BuildStats#

class yobx.container.BuildStats[source]#

Holds statistics collected during model export by ExtendedModelContainer.

Any key whose name begins with one of the allowed prefixes listed in _PREFIXES is accepted. Values default to 0.0 when first read via __getitem__().

Example:

stats = BuildStats()
stats["time_export_write_model"] += 0.5
stats["time_export_tobytes"] += 0.1

to_dict() → dict[str, float | int | bool | str][source]#: Returns the data in a dictionary, does not ùake a copy.

validate()[source]#: Validates that statistics startswith a known prefix.

ExportArtifact#

Standard output of every to_onnx() conversion function.

Every top-level to_onnx function (sklearn, tensorflow, litert, torch, sql …) returns an ExportArtifact instead of a bare ModelProto or ExtendedModelContainer. The instance bundles the exported proto, the optional large-model container, an ExportReport describing the export process, and an optional filename.

Parameters:

proto – ModelProto | FunctionProto | GraphProto | None The ONNX proto produced by the export. When large_model was requested the proto contains placeholders for external data; use get_proto() to obtain a fully self-contained proto.
container – ExtendedModelContainer | None The ExtendedModelContainer produced when the conversion was called with large_model=True. None otherwise.
report – ExportReport | None Statistics and metadata about the export.
filename – str | None Path where the model was last saved, or None if never saved.
builder – GraphBuilderExtendedProtocol Keeps the builder building the onnx model.
function – FunctionPieces | None When the export was performed with function_options.export_as_function=True, this holds a FunctionPieces instance with function-specific data (initializer names/values and nested local functions). None for regular model exports.

Example:

import numpy as np
from sklearn.linear_model import LinearRegression
from yobx.sklearn import to_onnx
from yobx.container import ExportArtifact, ExportReport

X = np.random.randn(20, 4).astype(np.float32)
y = X @ np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
reg = LinearRegression().fit(X, y)

artifact = to_onnx(reg, (X,))
assert isinstance(artifact, ExportArtifact)
assert isinstance(artifact.report, ExportReport)

proto = artifact.get_proto()
artifact.save("model.onnx")

SerializeToString() → bytes[source]#: Serializes the model to bytes. It does not includes weights if the model is stored in a container.

compute_node_stats(verbose: int = 0) → ExportArtifact[source]#

Compute per-op-type node counts and estimated FLOPs and store them in the report.

Delegates to ExportReport.compute_node_stats() using this artifact’s proto. Does nothing when proto is not a ModelProto.

Parameters:: verbose – verbosity level passed to ModelStatistics.
Returns:: self, to allow method chaining.

compute_symbolic_flops(verbose: int = 0) → ExportArtifact[source]#

Compute per-node symbolic FLOPs and store them in the report.

Delegates to ExportReport.compute_symbolic_flops() using this artifact’s proto. Does nothing when proto is not a ModelProto.

Parameters:: verbose – verbosity level passed to BasicShapeBuilder.
Returns:: self, to allow method chaining.

property functions: Sequence[FunctionProto]#: Returns the opset import.

get_proto(include_weights: bool = True) → Any[source]#

Return the ONNX proto, optionally with all weights inlined.

When the export was performed with large_model=True (i.e. container is set), the raw proto has external-data placeholders instead of embedded weight tensors. Passing include_weights=True (the default) uses to_ir() to build a fully self-contained ModelProto.

Parameters:: include_weights – when True (default) embed the large initializers stored in container into the returned proto. When False return the raw proto as-is.
Returns:: ModelProto, FunctionProto, or GraphProto.

Example:

artifact = to_onnx(estimator, (X,), large_model=True)
# Fully self-contained proto (weights embedded):
proto_with_weights = artifact.get_proto(include_weights=True)
# Proto with external-data placeholders:
proto_no_weights = artifact.get_proto(include_weights=False)

property graph: GraphProto#: Returns the GraphProto is the model is available. Fails otherwise.

property input_names: List[str]#: Returns the input names.

property ir_version: int#: Returns the opset import.

classmethod load(file_path: str, load_large_initializers: bool = True) → ExportArtifact[source]#

Load a saved model from file_path.

If the file references external data (i.e. the model was saved with large_model=True) an ExtendedModelContainer is created and returned in container. Otherwise the proto is loaded directly with onnx.load() and container is None.

Parameters:

file_path – path to the .onnx file.
load_large_initializers – when True (default) also load the large initializers stored alongside the model file.

Returns:

ExportArtifact with filename set to file_path.

Example:

artifact = ExportArtifact.load("model.onnx")
proto = artifact.get_proto()

property metadata_props: Sequence[StringStringEntryProto]#: Returns the opset import.

property opset_import: Sequence[OperatorSetIdProto]#: Returns the opset import.

property output_names: List[str]#: Returns the output names.

save(file_path: str, all_tensors_to_one_file: bool = True) → Any[source]#

Save the exported model to file_path.

When a ExtendedModelContainer is present (large_model=True was used during export) the model and its external weight files are saved via save(). Otherwise the proto is saved with onnx.save_model().

Parameters:

file_path – destination file path (including .onnx extension).
all_tensors_to_one_file – when saving a large model, write all external tensors into a single companion data file.

Returns:

the saved ModelProto.

Example:

artifact = to_onnx(estimator, (X,))
artifact.save("model.onnx")

save_report(onnx_path: str) → None[source]#: Save the report as an Excel file alongside onnx_path when available.

update(data: Any)[source]#: Updates report.

ExportReport#

class yobx.container.ExportReport(stats: List[Dict[str, Any]] | None = None, extra: Dict[str, Any] | None = None, build_stats: BuildStats | None = None, node_stats: List[Dict[str, Any]] | None = None, symbolic_flops: List[Dict[str, Any]] | None = None)[source]#

Holds statistics and metadata gathered during an ONNX export.

The _stats attribute stores the per-pattern optimization statistics returned by to_onnx() when called with return_optimize_report=True. Each element of the list is a dict with at least the keys "pattern", "added", "removed", and "time_in".

Additional arbitrary key-value pairs can be recorded via the update() method and are stored in extra.

Node-level statistics (op-type counts and estimated FLOPs) are stored in node_stats and can be populated by calling compute_node_stats().

Example:

report = ExportReport()
report.update({"time_total": 0.42})
print(report)

compute_node_stats(model: ModelProto, verbose: int = 0) → ExportReport[source]#

Compute per-op-type node counts and estimated FLOPs from model.

Delegates to model_statistics() for shape inference and FLOPs estimation. Results are stored in node_stats as a list of dicts with the keys:

op_type: ONNX operator type name (e.g. "MatMul").
count: Number of nodes with this op_type.
flops: Estimated total floating-point operations for all nodes of this type. None when the estimate is unavailable (shapes not fully known or the op is not covered by the estimator).

The rows are sorted by count descending.

Parameters:

model – exported ONNX model proto.
verbose – verbosity level passed to ModelStatistics.

Returns:

self, to allow method chaining.

<<<

import onnx.helper as oh
import onnx
from yobx.container import ExportReport

X = oh.make_tensor_value_info("X", onnx.TensorProto.FLOAT, [4, 8])
Y = oh.make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [4, 8])
node = oh.make_node("Relu", ["X"], ["Y"])
graph = oh.make_graph([node], "g", [X], [Y])
model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)])

report = ExportReport()
report.compute_node_stats(model)
print(report.node_stats)

>>>

    [{'op_type': 'Relu', 'count': 1, 'flops': 32}]

compute_symbolic_flops(model: ModelProto, verbose: int = 0) → ExportReport[source]#

Compute per-node symbolic FLOPs from model using symbolic shape inference.

Uses BasicShapeBuilder with inference=InferenceMode.COST to walk every node and record the FLOPs expression. When the model’s input shapes contain symbolic dimensions (strings such as "batch" or "seq"), the returned expressions are symbolic arithmetic strings (e.g. "2*batch*d_head*seq*seq"). For models with fully static shapes the values are plain integers. Operators whose cost cannot be estimated produce None.

Results are stored in symbolic_flops as a list of dicts with the keys:

op_type: ONNX operator type name (e.g. "MatMul").
node_name: The name of the ONNX node (may be an empty string).
symbolic_flops: Symbolic FLOPs expression (str), concrete integer, or None when the estimate is unavailable.

Parameters:

model – exported ONNX model proto.
verbose – verbosity level passed to BasicShapeBuilder.

Returns:

self, to allow method chaining.

<<<

import onnx.helper as oh
import onnx
from yobx.container import ExportReport

X = oh.make_tensor_value_info("X", onnx.TensorProto.FLOAT, ["batch", "d"])
Y = oh.make_tensor_value_info("Y", onnx.TensorProto.FLOAT, ["batch", "d"])
node = oh.make_node("Relu", ["X"], ["Y"])
graph = oh.make_graph([node], "g", [X], [Y])
model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)])

report = ExportReport()
report.compute_symbolic_flops(model)
print(report.symbolic_flops)

>>>

    [{'op_type': 'Relu', 'node_name': '', 'symbolic_flops': 'batch*d'}]

to_dict() → Dict[str, Any][source]#

Return a plain dictionary representation of this report.

Returns:: dictionary with keys "stats" and "extra".

to_excel(path: str) → None[source]#

Write the report contents to an Excel workbook at path.

The workbook contains up to six sheets:

stats: Every row from stats as a raw DataFrame. Only written when stats is non-empty.
stats_agg: The same data aggregated (summed) by "pattern", sorted by the number of removed nodes in descending order. Only written when stats is non-empty and contains a "pattern" column.
extra: extra rendered as a two-column table (key, value). Only written when extra is non-empty.
build_stats: build_stats rendered as a two-column table (key, value). Only written when build_stats is not None and contains at least one entry.
node_stats: Per-op-type node counts and estimated FLOPs from node_stats. Only written when node_stats is non-empty. Populated by compute_node_stats().
symbolic_flops: Per-node symbolic FLOPs expressions from symbolic_flops. Each row has op_type, node_name, and symbolic_flops columns. Only written when symbolic_flops is non-empty. Populated by compute_symbolic_flops().

Parameters:: path – destination file path (e.g. "report.xlsx").

<<<

import tempfile, os
from yobx.container import ExportReport

report = ExportReport(
    stats=[
        {"pattern": "p1", "added": 1, "removed": 2, "time_in": 0.01},
    ],
    extra={"time_total": 0.42},
)
with tempfile.TemporaryDirectory() as tmp:
    path = os.path.join(tmp, "report.xlsx")
    report.to_excel(path)
    print(f"Saved to {os.path.basename(path)!r}")

>>>

    Saved to 'report.xlsx'

to_string() → str[source]#

Return a human-readable text summary of this report.

The output includes:

Any extra key/value pairs stored in extra.
A per-pattern optimization table built from stats (when non-empty), grouped and sorted by the number of nodes removed.
Timing entries from build_stats (when available).

Returns:: multi-line string suitable for printing.

<<<

from yobx.container import ExportReport

report = ExportReport(
    stats=[
        {"pattern": "p1", "added": 1, "removed": 2, "time_in": 0.01},
        {"pattern": "p1", "added": 0, "removed": 1, "time_in": 0.02},
    ],
    extra={"time_total": 0.42},
)
print(report.to_string())

>>>

    -- extra --
      time_total: 0.42
    -- stats (aggregated by pattern) --
             added  removed  time_in
    pattern                         
    p1           1        3     0.03

update(data: Any)[source]#

Appends data to the report.

Parameters:: data – anything
Returns:: self

ExtendedModelContainer#

class yobx.container.ExtendedModelContainer(*args, **kwargs)[source]#

Overwrites onnx.model_container.ModelContainer to support torch tensors.

get_model_with_data()[source]#: Returns a copy of the model with data included in it.

get_prop(tensor: TensorProto) → StringStringEntryProto[source]#: Returns the location for a tensor stored in external data.

get_raw_data(np_tensor: ndarray | torch.Tensor | tf.Tensor) → bytes[source]#: Returns the raw data of a specific tensor whatever it is.

load(file_path: str, load_large_initializers: bool = True) → ExtendedModelContainer[source]#

Loads the large model.

Parameters:

file_path – model file
load_large_initializers – loads the large initializers, if not done, the model is incomplete but it can be used to look into the model without executing it and method onnx.model_container.ModelContainer._load_large_initializers can be used to load them later

Returns:

self

save(file_path: str, all_tensors_to_one_file: bool = True) → ModelProto[source]#

Saves the large model. The function returns a ModelProto, the current one if the model did not need any modification, a modified copy of it if it required changes such as giving file names to every external tensor.

Parameters:

file_path – model file
all_tensors_to_one_file – saves all large tensors in one file or one file per large tensor

Returns:

the saved ModelProto

to_ir() → onnx_ir.Model[source]#: Conversion to onnx_ir.Model.

yobx.container#

BuildStats#

ExportArtifact#

ExportReport#

ExtendedModelContainer#

This Page