[ディープラーニング] ONNX モデル ファイルはノードの名前を変更し、入力名を変更し、出力名を変更します

onnx モデル ファイルのノード名を変更する場合は、元の pytorch コード内で変更するか、onnx モデル ファイル内で直接変更する必要があります。

ここでは、onnx モデル ファイルで直接変更します。onnx ファイルがあり、出力ノードの名前は次のようになります:

変更しないでください。Python の変更を使用してください:

import onnx

# 加载模型文件
model = onnx.load("model.onnx")

# 修改输入输出张量的名称
idx_start = 0
for input in model.graph.input:
    for node in model.graph.node:
        # 如果当前节点的输入名称与待修改的名称相同，则将其替换为新名称
        for i, name in enumerate(node.input):
            if name == input.name:
                node.input[i] = "input_" + str(idx_start)
    input.name = "input_" + str(idx_start)
    idx_start += 1

idx_start = 0
for output in model.graph.output:
    for node in model.graph.node:
        # 如果当前节点的输入名称与待修改的名称相同，则将其替换为新名称
        for i, name in enumerate(node.output):
            if name == output.name:
                node.output[i] = "output_" + str(idx_start)
    output.name = "output_" + str(idx_start)
    idx_start += 1

# 保存修改后的模型
onnx.save(model, "modified_model.onnx")

変更後:

実は他のノードの名前も同様に変更することができますが、前後に接続されているノードに注意して変更する必要があります。

変更された内容が正しいことを確認します。

import numpy as np
import onnxruntime as ort

# 生成随机图像
# img = np.random.randint(-1, 1, (1, 3, 640, 640)).astype(np.float32)

img = np.load("img.npy")

# 加载模型
session = ort.InferenceSession("modified_model.onnx")
# 运行模型
output_new = session.run(None, {
    
    "input_0": img})

# 加载模型
session = ort.InferenceSession("model.onnx")
# 运行模型
output_old = session.run(None, {
    
    "x": img})

# 判断是否完全一样
print(np.allclose(output_new, output_old))