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 [source]#
- onnx.tools.net_drawer.GetOpNodeProducer(embed_docstring: bool = False, **kwargs: Any) Callable[[NodeProto, int], Node] [source]#
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 [source]#
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')