模型导出与ONNX转换：从PyTorch到嵌入式的第一步

说实话，模型量化部署这件事，很多人一上来就盯着量化算法调参。但我个人的经验是——导出这一步要是没走稳，后面全是白搭。ONNX就像一座桥，桥没搭好，再好的车也过不去。

今天咱们就把这座桥怎么搭、怎么验、怎么调，一次性说清楚。

为什么非要用ONNX？

PyTorch模型在PC上跑得欢，但NPU不认识它。NPU只认自己的一套指令集。ONNX的作用就是做一个中间表示——把PyTorch的计算图翻译成NPU能理解的语言。

我在项目中遇到过好几次，模型在PyTorch里精度99%，转成ONNX后直接崩了。原因五花八门：算子不支持、动态轴没设对、Opset版本太老……所以这一步，咱们得认真对待。

PyTorch模型转ONNX：torch.onnx.export 实战

导出代码其实就几行，但参数设置才是关键。我直接给一个我常用的模板：

import torch
import torch.onnx

# 加载你的YOLO模型
model = torch.load('yolov8n.pt', map_location='cpu')
model.eval()

# 构造一个虚拟输入
dummy_input = torch.randn(1, 3, 640, 640)

# 导出ONNX
torch.onnx.export(
    model,                     # 模型
    dummy_input,               # 输入张量
    'yolov8n.onnx',            # 输出文件名
    export_params=True,        # 导出参数
    opset_version=11,          # Opset版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],     # 输入名称
    output_names=['output'],   # 输出名称
    dynamic_axes={
        'input': {0: 'batch_size', 2: 'height', 3: 'width'},
        'output': {0: 'batch_size'}
    }
)
print("ONNX导出完成！")

这里有几个坑，我一个个说。

1. Opset版本：选对才能跑

Opset版本说白了就是ONNX的“方言版本”。版本越高，支持的算子越多，但NPU不一定认识。

Opset版本	特点	适用场景
10	基础算子，兼容性最好	老旧NPU、ARM CPU
11	支持更多切片、拼接操作	大部分NPU，推荐首选
12	支持GatherElements等	新NPU，需验证
13	支持Softmax优化	特定NPU加速库

2. 动态轴与静态轴：别让模型“僵住”

你想想看，训练时batch size可能是16，但部署时可能一次只处理1张图。如果导出时把batch size固定死了，推理时就只能按16来跑——浪费资源。

动态轴就是告诉ONNX：这个维度是可变的。我一般这样设：

dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size'}
}

但注意，不是所有NPU都支持动态轴。有些NPU要求所有维度固定，这时候就得用静态轴。我在一个项目里就吃过这个亏——NPU只支持固定640x640输入，我设了动态height/width，结果编译报错。后来改成静态才通过。

ONNX模型结构验证：导出不等于成功

导出完ONNX文件，你以为就完事了？我曾经有一次，导出很顺利，但加载到NPU上推理结果全是NaN。后来一查，是某个算子在ONNX里被替换成了不兼容的实现。

所以，验证这一步不能省。我一般做三层验证：

1. 语法验证：用onnx.checker检查模型结构是否合法
2. 形状验证：用onnx.shape_inference确认输入输出形状正确
3. 数值验证：用onnxruntime跑一遍，对比PyTorch输出

代码示例：

import onnx
import onnxruntime
import numpy as np

# 1. 语法验证
onnx_model = onnx.load('yolov8n.onnx')
onnx.checker.check_model(onnx_model)
print("语法验证通过！")

# 2. 形状验证
inferred_model = onnx.shape_inference.infer_shapes(onnx_model)
onnx.save(inferred_model, 'yolov8n_inferred.onnx')
print("形状验证通过！")

# 3. 数值验证
ort_session = onnxruntime.InferenceSession('yolov8n.onnx')
ort_inputs = {ort_session.get_inputs()[0].name: np.random.randn(1, 3, 640, 640).astype(np.float32)}
ort_outputs = ort_session.run(None, ort_inputs)
print(f"ONNX输出形状: {ort_outputs[0].shape}")
print("数值验证通过！")

Opset版本兼容性问题排查

这是最让人头疼的部分。你导出的ONNX在PC上跑得好好的，一上NPU就报错。常见的错误有：

• Unsupported Operator：NPU不认识某个ONNX算子
• Version Mismatch：Opset版本与NPU编译器不匹配
• Shape Inference Failed：动态轴导致形状推导失败

我的排查流程是这样的：

1. 先看错误日志，定位到具体算子
2. 查NPU的算子支持列表，看这个算子是否被支持
3. 如果不支持，尝试用onnx-simplifier简化模型
4. 还不行，就手动替换算子（比如把Softmax换成简化版本）
5. 最后一步：降Opset版本

举个例子，YOLO的NMS（非极大值抑制）在ONNX里是个复杂算子，很多NPU不支持。我一般会在导出前去掉NMS，把检测框输出交给后处理来做。

# 去掉NMS的导出方式
class YOLOWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    def forward(self, x):
        # 只输出原始检测结果，不做NMS
        return self.model.model(x)

wrapped_model = YOLOWrapper(model)
torch.onnx.export(
    wrapped_model,
    dummy_input,
    'yolov8n_no_nms.onnx',
    opset_version=11,
    # ...
)

本章知识结构

下面这张图总结了模型导出与ONNX转换的核心流程：

嗯，这张图把整个流程串起来了。从PyTorch模型开始，经过导出、验证，最后部署到NPU。每一步都有坑，但每一步也都有解法。

我个人习惯是：导出前先查NPU文档，导出后做三层验证，遇到问题先降Opset版本。这套流程帮我省了不少时间。