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

该函数的设计使得该函数可以重新应用于修改后的模型，并且可以替换 *ConstantOfShape* 为 *Range* 运算符，或者替换所有 *ConstantOfShape* 的填充值。