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[原始碼]¶

此函式會將模型的輸入和輸出維度大小更新為 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)[原始碼]¶

將初始值設定項或常數節點替換為 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 的填充值。