3、模型导出与ONNX转换：PyTorch模型转ONNX，解决动态轴与算子兼容性问题

说实话，模型训练完了只是第一步。真正要把模型塞进嵌入式设备，你绕不开ONNX这个中间格式。我刚开始做部署时，总觉得PyTorch模型直接转ONNX就是一行代码的事，结果被各种报错折磨得够呛。今天咱们就把这块硬骨头啃下来。

3.1 为什么非得转ONNX？

你想想看，PyTorch模型在PC上跑得飞起，但到了ARM Cortex-M或者NPU上，它根本不认识。ONNX就像个通用翻译器——把PyTorch的"方言"转成业界通用的"普通话"。我个人习惯是：所有模型先转ONNX，再转目标平台的格式。这样万一换平台，只需要改最后一步。

ONNX的好处很明显：

• 跨框架互通：PyTorch训练的模型，可以部署到TensorRT、OpenVINO、Tengine等推理引擎上
• 静态图优化：ONNX是静态图，推理引擎可以做常量折叠、算子融合等优化
• 量化支持：很多量化工具都基于ONNX做后训练量化

3.2 基础导出：一行代码搞定？

嗯，这里要注意。PyTorch官方提供了torch.onnx.export接口，但绝不是简单传个模型就行。我见过太多人直接复制官方示例，结果导出后推理结果全是错的。

先看一个最基础的导出示例：

import torch
import torch.onnx

# 假设你有一个训练好的轻量化语义分割模型
model = LightSegModel()
model.load_state_dict(torch.load('best_model.pth'))
model.eval()

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

# 导出ONNX
torch.onnx.export(
    model,                     # 模型
    dummy_input,               # 虚拟输入
    'lightseg.onnx',           # 输出文件名
    input_names=['input'],     # 输入节点名称
    output_names=['output'],   # 输出节点名称
    opset_version=11,          # 算子集版本
    do_constant_folding=True   # 常量折叠优化
)

print("ONNX导出成功！")

这段代码看着简单，但有几个坑。我刚开始做的时候，opset_version随便填了个13，结果板端推理引擎不支持，又得重导。建议先用11，兼容性最好。

3.3 动态轴：输入尺寸不固定怎么办？

语义分割模型有个特点——输入图像尺寸可能变化。比如你训练时用256x256，但实际部署时可能是320x240或者别的分辨率。这时候就需要动态轴。

为什么会这样？因为ONNX默认把所有维度都固定死了。你导出时传了(1, 3, 256, 256)，它就认为输入永远是这个形状。推理时换个尺寸，直接报错。

解决办法是显式指定哪些维度是动态的：

torch.onnx.export(
    model,
    dummy_input,
    'lightseg_dynamic.onnx',
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size', 2: 'height', 3: 'width'},
        'output': {0: 'batch_size', 2: 'height', 3: 'width'}
    },
    opset_version=11
)

这里dynamic_axes字典的key是节点名称，value是维度索引和别名。我一般把batch、height、width都设成动态，这样部署时灵活很多。

3.4 算子兼容性：最常见的翻车现场

算子兼容性问题，说白了就是PyTorch里有的操作，ONNX不支持。我统计过，80%的导出失败都是这个原因。

常见的"问题算子"包括：

PyTorch算子	问题描述	解决方案
torch.einsum	ONNX不支持爱因斯坦求和	拆成reshape+matmul
F.adaptive_avg_pool2d	输出尺寸为1时没问题，动态尺寸可能报错	用avg_pool2d配合固定kernel
torch.topk	ONNX opset 11以下不支持	升级opset或用sort替代
nn.Dropout	推理模式下没问题，但训练模式导出会报错	确保model.eval()
torch.cat + 动态轴	某些推理引擎对cat的维度有限制	尽量在固定维度上cat

我记得有一次，模型里用了torch.meshgrid做坐标映射，ONNX死活不支持。最后我手动写了个循环替代，虽然效率低点，但至少能导出。

3.5 验证导出的ONNX模型

导出完别急着部署，先验证一下。我习惯用ONNX Runtime做本地推理，对比PyTorch和ONNX的输出差异：

import onnx
import onnxruntime as ort
import numpy as np

# 1. 检查ONNX模型结构
onnx_model = onnx.load('lightseg.onnx')
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过！")

# 2. 用ONNX Runtime推理
ort_session = ort.InferenceSession('lightseg.onnx')
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_outputs = ort_session.run(None, ort_inputs)

# 3. 对比PyTorch输出
with torch.no_grad():
    torch_output = model(dummy_input)

# 计算误差
diff = np.abs(ort_outputs[0] - torch_output.numpy())
print(f"最大误差: {diff.max():.6f}")
print(f"平均误差: {diff.mean():.6f}")

一般来说，误差在1e-5以内就算正常。如果误差很大，说明导出过程中某些算子被替换了，需要排查。

3.6 轻量化模型的特殊处理

对于轻量化语义分割模型（比如MobileNetV3+LRASPP、STDC、PIDNet等），有几个额外注意点：

• 深度可分离卷积：ONNX对nn.Conv2d的groups参数支持良好，但某些推理引擎对groups>1的卷积优化不够。我建议导出后检查一下ONNX图，看看有没有被拆成多个普通卷积。
• SE模块：Squeeze-and-Excitation里的sigmoid和mul操作，ONNX支持没问题。但如果你用了hard_sigmoid，记得用relu6(x+3)/6替代，兼容性更好。
• 上采样层：语义分割模型最后通常有上采样。我建议用nn.Upsample(scale_factor=4, mode='bilinear', align_corners=False)，不要用size参数，否则动态尺寸导出会报错。

3.7 实战经验总结

最后分享几个我踩过的坑：

1. 导出前一定要model.eval()。我刚开始忘了这步，结果导出的ONNX里还带着Dropout，推理结果随机变化。
2. 虚拟输入的shape要和训练时一致。如果你训练时用了torch.cat或者torch.add，输入尺寸变了可能导致shape推导失败。
3. opset_version不是越高越好。高版本支持更多算子，但嵌入式推理引擎可能不支持。我一般用11，兼顾功能和兼容性。
4. 用onnx-simplifier简化模型。导出后的ONNX可能包含很多冗余节点，用python -m onnxsim lightseg.onnx lightseg_sim.onnx可以去掉这些冗余，减少模型体积。

嗯，模型导出这块就讲这么多。说白了，ONNX转换就是个"翻译"过程，翻译得对不对，得靠验证说话。下一章咱们聊聊如何用Tengine在ARM板上跑起来这个ONNX模型。