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 [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]¶
此函数将模型的输入和输出的维度大小更新为 input_dims 和 output_dims 中提供的值。如果提供的维度值为负数,则将为此维度设置一个唯一的 dim_param。
示例:如果输入和输出的形状如下
shape(input_1) = (‘b’, 3, ‘w’, ‘h’)
shape(input_2) = (‘b’, 4)
shape(output) = (‘b’, ‘d’, 5)
参数可以按如下方式提供
input_dims = { "input_1": ['b', 3, 'w', 'h'], "input_2": ['b', 4], } output_dims = { "output": ['b', -1, 5] }
综合起来
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)[source]¶
将初始化器或常量节点替换为 ConstantOfShape 节点以减小大小。
这降低了编写针对特定图结构的单元测试的成本。
- 参数:
onx – ModelProto
threshold – 小于此阈值的每个初始化器不受影响
ir_version – 对于 ir_version <= 3,初始化器必须指定为输入,如果 onx 是
FunctionProto
或GraphProto
,则必须指定此参数use_range – 如果使用 Range 算子代替 ConstantOfShape 以避免常量张量
value_constant_of_shape – 用作所有 ConstantOfShape 节点的填充值,较高的值可能会产生 nan 或 inf 预测结果
- 返回值:
onx,修改后的 ModelProto
设计此函数的目的是使其可以在修改后的模型上重新应用,并可以用 Range 算子替换 ConstantOfShape,或替换每个 ConstantOfShape 的填充值。