graphsanitizer

Graph sanitization and optimization for ONNX models.

Classes

GraphSanitizer

A class for sanitizing ONNX model graphs, a part of the AutoCast tool.
class GraphSanitizer

Bases: object

A class for sanitizing ONNX model graphs, a part of the AutoCast tool.

__init__(model, min_opset=13, max_ir_version=None, trt_plugins=[])

Initialize GraphSanitizer.

Parameters:

  • model (ModelProto) – ONNX model to sanitize

  • min_opset (int) – minimum opset version to use

  • max_ir_version (int | None) – maximum IR version supported by ORT

  • trt_plugins (list[str] | None) – list of TensorRT plugin library paths in .so format (compiled shared library).

Return type:

None

cleanup_model()

Use GraphSurgeon to cleanup unused nodes, tensors and initializers.

Return type:

None

convert_opset()

Convert the model to the given opset version.

Parameters:

min_opset – minimum opset version to use

Return type:

None

The method checks all opset imports and converts the model if any are below the minimum version.

ensure_graph_name_exists()

Ensures that the model’s name exists.

Return type:

None

find_custom_nodes()

Find custom nodes in the model.

Scans through all nodes in the graph and logs any nodes that use custom operators that are not part of the standard ONNX operator set.

Return type:

None

remove_disconnected_outputs()

Remove disconnected outputs from the model.

Return type:

None

replace_custom_domain_nodes()

Replace custom domain nodes with standard ONNX nodes.

replace_layernorm_pattern()

Detects and replaces LayerNorm operation patterns.

This method scans through the graph looking for sequences of operations that implement LayerNorm functionality and replaces them with the more efficient LayerNormalization operator.

Return type:

None

sanitize()

Sanitize the model graph.

Currently, this finds decomposed LayerNorm patterns and replaces them with a single LayerNormalization operator. Additional functionality may be added in the future.

Return type:

None

set_ir_version(max_ir_version)

Set the model’s IR version to the maximum supported version.

Parameters:

max_ir_version (int | None) – maximum IR version to use.

Return type:

None

The method checks the IR version and cuts it off at the maximum supported version. See https://onnxruntime.ai/docs/reference/compatibility.html#onnx-opset-support