4. 模型导出与转换：PyTorch 模型转 ONNX 的常见错误、ONNX 转 TensorRT/NCNN/MNN 的踩坑记录、动态轴与静态轴设置

模型导出与转换，说白了就是让模型换个“语言”去跟硬件打交道。PyTorch 训练好的模型再牛，到了边缘设备上也得“入乡随俗”。我这些年踩过的坑，十个手指头都数不过来。今天就把这些血泪史整理成一份避坑指南，希望能帮你少走弯路。

4.1 PyTorch 转 ONNX：那些让你抓狂的常见错误

ONNX 是个中间格式，像个“翻译官”。但 PyTorch 和 ONNX 之间，经常“翻译”出问题。

4.1.1 算子不支持：最头疼的问题

PyTorch 里有些骚操作，ONNX 根本不认识。比如 torch.einsum、torch.where 的某些用法，或者自定义的 torch.autograd.Function。

解决办法其实就几种：

• 替换算子：把不支持的算子换成 ONNX 支持的等价实现。比如用 torch.sort 加切片代替 torch.topk。
• 自定义算子：实在绕不过去，就写个自定义 ONNX 算子。但我不建议新手这么干，调试起来很痛苦。
• 升级版本：ONNX 和 PyTorch 的版本匹配很重要。我习惯用 PyTorch 1.13 + ONNX 1.14，兼容性相对好一些。

4.1.2 动态轴与静态轴：一个参数搞死你

这是新手最容易忽略的地方。ONNX 导出时，dynamic_axes 参数决定了哪些维度是动态的。

# 错误示范：没设动态轴
torch.onnx.export(model, dummy_input, "model.onnx")

# 正确做法：明确指定动态轴
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    dynamic_axes={
        "input": {0: "batch_size", 2: "height", 3: "width"},
        "output": {0: "batch_size"}
    },
    opset_version=14
)

你想想看，如果没设动态轴，导出的 ONNX 模型就固定死了输入尺寸。部署到边缘设备上，换个分辨率就得重新导出，多麻烦。

4.2 ONNX 转 TensorRT：精度与速度的博弈

TensorRT 是 NVIDIA 的优化利器，但转换过程经常出幺蛾子。

4.2.1 FP16 精度损失：肉眼可见的退化

转 TensorRT 时，很多人直接开 FP16 加速。结果模型跑起来快了，但输出结果完全不对。

为什么会这样？因为 MobileViT 里有些层对精度特别敏感，比如 LayerNorm 和 Softmax。FP16 下这些层的计算误差会被放大。

4.2.2 动态 shape 的坑：TensorRT 的“静态癖”

TensorRT 对动态 shape 的支持很有限。虽然从 TensorRT 8.0 开始有了 optimization profile，但用起来还是各种限制。

# TensorRT 动态 shape 配置示例
profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 3, 224, 224), (1, 3, 256, 256), (1, 3, 512, 512))
config.add_optimization_profile(profile)

这里要注意：min、opt、max 三个 shape 必须设置合理。我见过有人把 min 设成 (1, 3, 1, 1)，结果推理时直接崩了。因为 MobileViT 的 patch embedding 层对输入尺寸有最小要求。

4.3 ONNX 转 NCNN：移动端的“瘦身”挑战

NCNN 是腾讯开源的移动端推理框架，对 ARM 平台优化得很好。但 MobileViT 这种带 Transformer 的模型，转 NCNN 时问题不少。

4.3.1 不支持的操作：GELU 激活函数

MobileViT 用了 GELU 激活函数，但 NCNN 原生不支持。我当初查了半天文档，发现得用 GELU 的近似实现来替换。

替换方法：在 PyTorch 导出 ONNX 之前，就把 GELU 换成 NCNN 支持的近似实现。或者在 ONNX 模型里手动修改节点。

4.3.2 内存优化：MobileViT 的“大胃口”

MobileViT 虽然叫“Mobile”，但 Transformer 结构的内存占用可不小。NCNN 默认的内存分配策略可能不够用。

我的经验是：

• 开启 NCNN 的 use_vulkan_compute，利用 GPU 加速。
• 调整 opt 参数里的 num_threads，别开太多线程，否则内存会爆。
• 对 MobileViT 的 attention 层做“内存复用”，减少临时张量的分配。

4.4 ONNX 转 MNN：阿里系的“快”与“坑”

MNN 是阿里巴巴开源的推理框架，主打轻量和高效。但 MobileViT 转 MNN 时，有几个点特别容易翻车。

4.4.1 维度顺序：NCHW vs NHWC

MNN 默认用 NHWC 格式，而 PyTorch 是 NCHW。转换时如果没处理好维度顺序，模型推理结果就是错的。

# MNN 转换时指定输入格式
mnnconvert --modelFile model.onnx --MNNModel model.mnn --inputDataFormat NHWC

但要注意：MobileViT 内部有些层对维度顺序有要求。我建议在 ONNX 导出时就固定为 NCHW，然后在 MNN 推理时做一次维度转换。

4.4.2 量化陷阱：INT8 的精度灾难

MNN 支持 INT8 量化，但 MobileViT 这种模型量化后精度下降很明显。我做过测试，INT8 量化后 Top-1 准确率掉了 3-5 个点。

4.5 知识体系总览

下面这张图总结了模型导出与转换的核心流程和关键节点。你可以把它当作一张“作战地图”。

4.6 总结：我的“三板斧”

说了这么多，其实核心就三点：

1. 版本匹配：PyTorch、ONNX、TensorRT/NCNN/MNN 的版本一定要对应好。我一般会在项目开始前建一个“版本兼容性表”，省得后面出问题。
2. 精度验证：每转一次，都要用同样的输入跑一遍，对比输出。误差超过 1% 就要排查原因。
3. 渐进式优化：别一上来就开 FP16 或 INT8。先跑通 FP32，再逐步优化。每一步都验证，出了问题也好定位。

嗯，模型导出与转换这块，说白了就是个“细心活”。你只要把每个环节的坑都摸清楚了，后面部署起来就顺风顺水。

可可精品资料库