4. 模型导出与转换

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

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

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

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

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

我曾经踩过的坑： 有一次我用 torch.topk 做特征筛选，导出 ONNX 时直接报错。后来发现 ONNX 的 TopK 算子只支持固定 k 值，而我传了个变量进去。

解决办法其实就几种：

  • 替换算子：把不支持的算子换成 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 模型就固定死了输入尺寸。部署到边缘设备上，换个分辨率就得重新导出，多麻烦。

我的个人习惯： 只要模型可能用在不同的输入尺寸上，我一定设动态轴。但要注意，动态轴太多会影响推理速度。我一般只把 batch size 和 H/W 设为动态，channel 保持静态。

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

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

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

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

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

我建议的做法：
  1. 先用 FP32 转一遍，验证精度没问题。
  2. 再开 FP16，但要对敏感层做“精度保护”。
  3. 用 TensorRT 的 setLayerPrecision 接口，把 LayerNorm 和 Softmax 强制设为 FP32。

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 的近似实现来替换。

我曾经踩过的坑： 直接硬转，结果 NCNN 推理时输出全是 NaN。后来我把 GELU 替换成了 x * sigmoid(1.702 * x) 的近似版本，才跑通。

替换方法：在 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 个点。

我的建议： 如果对精度要求高，先用 FP16 量化试试。INT8 量化需要做校准集，而且校准集的数据分布要跟实际场景一致。我一般会准备 500-1000 张跟部署场景相似的图片做校准。

4.5 知识体系总览

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

模型导出与转换核心流程 PyTorch 模型 ONNX 中间格式 TensorRT NCNN MNN 常见踩坑点 • 算子不支持：GELU、einsum、topk 等 • 动态轴设置：batch、H/W 动态，channel 静态 • 精度损失：FP16 下 LayerNorm、Softmax 敏感 • 维度顺序：NCHW vs NHWC 转换 • 内存优化：Transformer 结构内存占用大

4.6 总结：我的“三板斧”

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

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

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

一句话总结： 导出 ONNX 时设好动态轴，转 TensorRT 时保护敏感层精度，转 NCNN 时替换不支持的算子，转 MNN 时注意维度顺序和量化校准。每一步都验证，别偷懒。

可可精品资料库

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