3. 模型导出基础：PyTorch模型转ONNX

模型训练好了，怎么把它塞进嵌入式设备？

答案就是——先转成ONNX。ONNX就像模型界的通用语言。不管你是PyTorch、TensorFlow还是飞桨，都能转成ONNX。然后ONNX再转成各芯片的专有格式。这一步，是端侧移植的咽喉。

3.1 PyTorch模型转ONNX：其实就几行代码

PyTorch官方提供了torch.onnx.export接口。我个人习惯把它封装成一个函数，方便反复调用。

import torch
import torch.onnx

def export_to_onnx(model, dummy_input, onnx_path, opset_version=11):
    model.eval()
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        opset_version=opset_version,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes=None  # 先不设动态轴
    )
    print(f'ONNX模型已保存至: {onnx_path}')

# 使用示例
model = YourBEVModel()
dummy = torch.randn(1, 3, 224, 224)
export_to_onnx(model, dummy, 'bev_model.onnx')

这里有个关键参数——opset_version。它决定了ONNX算子集的版本。版本越高，支持的算子越多，但兼容性可能下降。我建议先设成11，这是目前最稳妥的版本。

3.2 ONNX算子兼容性检查：别让模型卡在半路

模型导出成功，不代表万事大吉。你想想看，PyTorch里有些算子，ONNX根本不认识。比如torch.einsum、torch.nn.functional.grid_sample（BEV里常用）。

怎么办？

第一步，用onnx.checker做基础检查。

import onnx

onnx_model = onnx.load('bev_model.onnx')
onnx.checker.check_model(onnx_model)
print('ONNX模型结构检查通过')

第二步，用onnxruntime做推理验证。这一步能发现运行时错误。

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession('bev_model.onnx')
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = sess.run([output_name], {input_name: dummy_input})
print('ONNX推理成功，输出shape:', output[0].shape)

第三步，也是最关键的——算子兼容性清单。我整理了一份常用算子对照表：

PyTorch算子	ONNX支持情况	常见问题
Conv2d / BatchNorm2d	✅ 完全支持	无
ReLU / Sigmoid / Softmax	✅ 完全支持	无
Add / Mul / Cat	✅ 完全支持	无
Reshape / Transpose	⚠️ 部分支持	动态shape可能失败
grid_sample	❌ 不支持	需用自定义算子或替换
einsum	❌ 不支持	需拆解为矩阵乘
MultiHeadAttention	⚠️ 部分支持	需拆解为基本算子

3.3 动态轴与静态轴设置：灵活还是高效？

ONNX模型里，每个张量都有shape。比如(1, 3, 224, 224)。这里的1是batch size。如果固定死，模型就只能处理batch=1的输入。这就是静态轴。

如果想让模型支持任意batch size，就得设动态轴。

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

设了动态轴，ONNX模型就能接受(N, 3, 224, 224)的输入。但代价是什么？

• 推理速度：动态shape会触发算子重新编译，速度下降10%~30%
• 内存占用：无法预分配固定内存，碎片化增加
• 芯片兼容性：很多端侧芯片不支持动态shape，直接报错

所以我的建议是：

• 量产部署：用静态轴，固定batch=1。性能最优。
• 调试阶段：用动态轴，方便测试不同输入。
• 特殊场景：如果必须支持多batch，先确认芯片工具链是否支持。

3.4 本章知识体系

下面这张图，帮你理清模型导出的完整流程：

从PyTorch到芯片工具链，中间三步缺一不可。算子兼容性检查帮你提前发现坑，动态轴设置决定部署灵活性，ONNX Runtime验证确保推理正确。这三步走完，模型导出才算真正完成。