HeadCountAlgo/demo/models/utils/output2output0.py

import onnx

# 1. 加载原始的ONNX模型
model = onnx.load("../yolo-crowd-ft-e60-grid.onnx")  # 请将 "原始模型.onnx" 替换为你的模型文件名

# 2. 修改节点输出名称
for node in model.graph.node:
    for i, output_name in enumerate(node.output):
        if output_name == "output":
            node.output[i] = "output0"
            print(f"节点 '{node.name}' 的输出名称已从 'output' 修改为 'output0'")

# 3. 修改模型输出名称
for output in model.graph.output:
    if output.name == "output":
        output.name = "output0"
        print("模型输出名称已从 'output' 修改为 'output0'")

# 4. 保存修改后的模型
onnx.save(model, "yolo-crowd-ft-e60.onnx")  # 可以指定新的文件名
print("模型已保存为 'yolo-crowd-ft-e60.onnx'")

# 5. 验证模型的正确性
try:
    onnx.checker.check_model(model)
    print("模型验证成功，没有发现错误。")
except onnx.checker.ValidationError as e:
    print(f"模型验证失败: {e}")