onnx.tools

net_drawer

onnx.tools.net_drawer.GetPydotGraph(graph: GraphProto, name: Optional[str] = None, rankdir: str = 'LR', node_producer: Optional[Callable[[NodeProto, int], Node]] = 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')