3、ONNX模型导出：PyTorch模型转ONNX、动态轴与静态轴设置、算子兼容性检查

好，咱们进入第三个实战环节。前面两章我们把Occupancy网络的理论和PyTorch实现捋清楚了，现在要开始干一件很实际的事——把模型从PyTorch搬到ONNX。

为什么要做这一步？说白了，嵌入式设备上很少能直接跑PyTorch。ONNX就像个中间人，帮我们把模型翻译成各种推理引擎都能看懂的语言。我最早做部署时，觉得这步就是走个过场，结果被各种坑折磨得够呛。嗯，今天咱们就把这些坑提前填上。

3.1 PyTorch转ONNX：基础操作

PyTorch官方提供了torch.onnx.export这个接口。用法其实不复杂，但有几个关键参数你得搞明白。

import torch
import torch.onnx

# 假设我们有一个训练好的Occupancy网络
model = OccupancyNetwork()
model.eval()  # 别忘了切到eval模式

# 创建一个虚拟输入
dummy_input = torch.randn(1, 3, 128, 128)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "occupancy_network.onnx",
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里我重点说几个容易踩坑的地方。

第一，opset_version怎么选？我个人习惯用11或12。太低的版本有些算子不支持，太高的版本嵌入式推理引擎可能还没跟上。我在项目中遇到过，选了opset 15导出的模型，结果在某个老版本的TensorRT上直接报错。所以我的建议是：先查一下你的目标推理引擎支持到哪个opset，再决定。

第二，do_constant_folding建议打开。这个参数会把一些固定的计算提前算好，比如权重和偏置的乘法。模型会变小，推理也会快一点。但要注意，如果你的模型里有动态形状相关的操作，可能会出问题。

小技巧：导出后可以用Netron工具打开ONNX文件，可视化检查网络结构。我每次导出都会看一眼，确认输入输出形状对不对，有没有多余的节点。

3.2 动态轴与静态轴：你选对了吗？

这是很多人容易搞混的地方。什么叫动态轴？说白了就是允许输入的形状在推理时变化。比如你的Occupancy网络，训练时用的batch size是4，但部署时可能一次只推理1张图，或者有时候需要同时处理8张图。

静态轴就是固定死的。比如输入必须是[1, 3, 128, 128]，多一张少一张都不行。好处是推理引擎可以做更多优化，速度更快。

动态轴则灵活得多。你可以在导出时指定哪些维度是可变的。

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

你看，这里我把batch、height、width都设成了动态。但要注意，不是所有算子都支持动态轴。比如一些reshape操作，如果形状是动态的，导出时可能会报错。

我个人经验是：能静态就别动态。我在一个项目中，为了图省事把所有轴都设成动态，结果推理速度比静态版本慢了将近一倍。后来我改成只把batch size设成动态，其他轴固定，速度就上来了。

注意：如果你的Occupancy网络输出是固定尺寸的（比如128x128），那height和width其实没必要设成动态。设成动态反而可能引入不必要的复杂度。

3.3 算子兼容性检查：别等部署了才发现

这个环节我吃过不少亏。你想想看，PyTorch里跑得好好的模型，导出成ONNX后，某些算子可能不被支持，或者被替换成了效率很低的实现。

怎么检查？我一般分三步走：

  1. 用ONNX Runtime验证：这是最直接的方法。导出后用ONNX Runtime跑一遍推理，对比PyTorch的输出。
  2. 检查算子支持列表：去ONNX官方文档查一下你的目标opset版本支持哪些算子。
  3. 用工具自动检查：比如onnx.checker和onnx.shape_inference。
import onnx
import onnxruntime as ort
import numpy as np

# 加载ONNX模型
onnx_model = onnx.load("occupancy_network.onnx")

# 检查模型结构是否合法
onnx.checker.check_model(onnx_model)

# 用ONNX Runtime推理
ort_session = ort.InferenceSession("occupancy_network.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: np.random.randn(1, 3, 128, 128).astype(np.float32)}
ort_outputs = ort_session.run(None, ort_inputs)

# 对比PyTorch输出
with torch.no_grad():
    torch_outputs = model(torch.from_numpy(ort_inputs['input']))

# 检查误差
diff = np.abs(ort_outputs[0] - torch_outputs.numpy())
print(f"最大误差: {diff.max()}")

如果误差在1e-5以内，基本没问题。如果误差很大，那就要排查了。我曾经遇到过一个情况，PyTorch里的F.interpolate在ONNX里被映射成了Resize算子，但坐标对齐方式不一样，导致输出差了0.1。后来我手动指定了coordinate_transformation_mode才解决。

常见不兼容算子：
  • torch.where — 有条件分支，容易出问题
  • torch.topk — 某些opset版本不支持
  • torch.nonzero — 输出形状不固定，动态轴下容易报错
  • F.grid_sample — Occupancy网络里常用，但ONNX支持有限

3.4 实战中的避坑指南

最后分享几个我实际踩过的坑，希望能帮你少走弯路。

坑一：Batch Normalization在eval模式下没冻结。我刚开始做部署时，忘了调用model.eval()，结果导出的ONNX里BN层还在用训练时的统计方式，推理结果完全不对。记住，导出前一定要切到eval模式。

坑二：输入输出的dtype不匹配。PyTorch默认是float32，但有些嵌入式推理引擎只支持float16。你可以在导出时用torch.onnx.export的dtype参数指定，或者导出后用onnx工具转换。

坑三：动态轴导致的内存爆炸。有一次我把height和width也设成动态，结果推理引擎为了支持任意尺寸，预分配了超大内存，直接OOM了。后来我限制了输入尺寸范围，才解决。

我的习惯：每次导出后，我都会用ONNX Runtime跑一遍完整的推理流程，包括前处理和后处理。这样能尽早发现问题，而不是等到部署到设备上才抓瞎。

好了，这一章的内容就到这里。ONNX导出看起来简单，但细节不少。下一章我们会讲如何用TensorRT对ONNX模型做量化加速，那才是真正让模型在嵌入式设备上跑起来的关键一步。

← 上一章 📖 返回目录 下一章 →