onnx.tools

net_drawer

onnx.tools.net_drawer.GetPydotGraph(graph: GraphProto, name: str | None = None, rankdir: str = 'LR', node_producer: Callable[[NodeProto, int], Node] | None = None, embed_docstring: bool = False) → Dot

onnx.tools.net_drawer.GetOpNodeProducer(embed_docstring: bool = False, **kwargs: Any) → Callable[[NodeProto, int], Node]

from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer

pydot_graph = GetPydotGraph(
    model_onnx.graph,  # model_onnx is a ModelProto instance
    name=model_onnx.graph.name,
    rankdir="TP",
    node_producer=GetOpNodeProducer("docstring"))
pydot_graph.write_dot("graph.dot")

update_inputs_outputs_dims

onnx.tools.update_model_dims.update_inputs_outputs_dims(model: ModelProto, input_dims: dict[str, list[Any]], output_dims: dict[str, list[Any]]) → ModelProto

This function updates the dimension sizes of the model’s inputs and outputs to the values provided in input_dims and output_dims. if the dim value provided is negative, a unique dim_param will be set for that dimension.

Example. if we have the following shape for inputs and outputs:

• shape(input_1) = (‘b’, 3, ‘w’, ‘h’)

• shape(input_2) = (‘b’, 4)

• shape(output) = (‘b’, ‘d’, 5)

The parameters can be provided as:

input_dims = {
    "input_1": ['b', 3, 'w', 'h'],
    "input_2": ['b', 4],
}
output_dims = {
    "output": ['b', -1, 5]
}

Putting it together:

model = onnx.load('model.onnx')
updated_model = update_inputs_outputs_dims(model, input_dims, output_dims)
onnx.save(updated_model, 'model.onnx')

replace_initializer_by_constant_of_shape

onnx.tools.replace_constants.replace_initializer_by_constant_of_shape(onx: FunctionProto | GraphProto | ModelProto, threshold: int = 128, ir_version: int | None = None, use_range: bool = False, value_constant_of_shape: float = 0.5)

Replace initializers or constant node by nodes ConstantOfShape to reduce the size.

This reduce the cost to write a unit test about a specific graph structure.

Parameters:

• onx – ModelProto

• threshold – every initializer under this threshold is not impacted

• ir_version – initializer must be specified as input for ir_version <= 3, this must be specified if onx is FunctionProto or GraphProto

• use_range – if uses operator Range instead of ConstantOfShape to avoid constant tensors

• value_constant_of_shape – value to use as a value for all nodes ConstantOfShape, a high value may produce nan or inf predictions

Returns:

onx, modified ModelProto

The function is designed so that the function can be reapplied on a modified model and either replace ConstantOfShape with Range operators, either replace the fill value for every ConstantOfShape.