2. ONNX模型优化与简化

任务目标

对给定的ONNX模型执行图优化：使用onnx-simplifier消除冗余运算（如常量折叠、算子融合），使用onnxoptimizer进行通用图优化（如公共子表达式消除、死节点删除），最终输出优化后的模型，并对比节点数量与文件大小。

环境依赖

pip install onnx onnxoptimizer onnx-simplifier

关键完整代码

import onnx
import onnxoptimizer
from onnxsim import simplify
import os

# 1. 加载原始模型（示例：构造一个含冗余恒等映射的模型）
import numpy as np
from onnx import helper, TensorProto

# 构造演示模型：两个连续Identity + 一个可常量折叠的Add
x = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 3, 4, 4])
y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 3, 4, 4])

const_val = np.random.randn(1, 3, 4, 4).astype(np.float32)
const_init = helper.make_tensor('const', TensorProto.FLOAT, [1, 3, 4, 4], const_val.flatten())

node1 = helper.make_node('Identity', ['input'], ['id1'])
node2 = helper.make_node('Identity', ['id1'], ['id2'])
node3 = helper.make_node('Add', ['id2', 'const'], ['output'])

graph = helper.make_graph([node1, node2, node3], 'demo_graph', [x], [y], initializer=[const_init])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
model.ir_version = 8

onnx.save(model, 'original.onnx')
print(f"原始模型节点数: {len(model.graph.node)}")

# 2. onnx-simplifier 优化（常量折叠 + 冗余节点消除）
sim_model, check = simplify(model)
assert check, "simplifier 验证失败"
onnx.save(sim_model, 'simplified.onnx')
print(f"simplifier 后节点数: {len(sim_model.graph.node)}")

# 3. onnxoptimizer 进一步优化（公共子表达式消除、死节点删除等）
optimized_model = onnxoptimizer.optimize(sim_model)
onnx.save(optimized_model, 'optimized.onnx')
print(f"optimizer 后节点数: {len(optimized_model.graph.node)}")

# 4. 验证：输出形状与数值一致性（使用onnxruntime）
import onnxruntime as ort

def run_inference(model_path):
    sess = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
    input_data = np.random.randn(1, 3, 4, 4).astype(np.float32)
    return sess.run(None, {'input': input_data})[0]

orig_out = run_inference('original.onnx')
opt_out = run_inference('optimized.onnx')
print(f"数值一致性 (最大误差): {np.max(np.abs(orig_out - opt_out)):.2e}")

# 5. 文件大小对比
for f in ['original.onnx', 'simplified.onnx', 'optimized.onnx']:
    print(f"{f}: {os.path.getsize(f)} bytes")

运行与验证

1. 将上述代码保存为optimize_onnx.py并执行：python optimize_onnx.py
2. 预期输出（确定性结果）：
   原始模型节点数: 3
simplifier 后节点数: 1（两个Identity被删除，Add被常量折叠为单一输出）
   optimizer 后节点数: 1（无进一步可消除节点）
   数值一致性 (最大误差): 0.00e+00
   文件大小：optimized.onnx 显著小于 original.onnx
3. 判定通过条件：
   - 优化后节点数少于原始节点数；
   - 优化前后推理输出的最大误差 < 1e-6；
   - 三个模型文件均成功生成且可被onnxruntime正常加载。

关键说明

• 示例模型故意构造了连续Identity和可常量折叠的Add，用于演示两类优化器的实际效果。真实模型只需替换onnx.load('your_model.onnx')即可复用相同流程。
• simplify()返回的布尔值表示优化后模型与原始模型在随机输入下的数值一致性验证是否通过；若为False，需检查模型是否包含动态形状或非标准算子。
• onnxoptimizer.optimize()默认启用全部安全优化pass，若需针对性调整，可传入pass列表参数（如['eliminate_deadend']）。