第3章 模型适配基础：PyTorch → ONNX → OM
模型适配，说白了就是把你在 GPU 上训好的模型，搬到昇腾设备上跑起来。这一步绕不开，也急不来。我刚开始接触昇腾平台时，觉得不就是换个后端吗？结果踩了不少坑。这一章咱们就把最核心的链路讲透：PyTorch 转 ONNX，ONNX 再转昇腾离线模型（OM），以及 ATC 工具怎么用。
3.1 为什么需要 ONNX 这个中间人？
你想想看，PyTorch 有自己的计算图表示，昇腾芯片又不直接认识它。ONNX 就像个通用翻译器——把 PyTorch 的「方言」转成业界标准中间格式。我个人习惯把 ONNX 看作模型交换的「世界语」，不管训练框架是什么，只要导出 ONNX，后续的部署工具都能接住。
我在项目中遇到过，有人想跳过 ONNX 直接转 OM，结果折腾了两天没跑通。后来规规矩矩走 ONNX 链路，半小时搞定。所以别嫌多这一步，它其实是帮你省时间。
3.2 PyTorch 模型转 ONNX
PyTorch 官方提供了 torch.onnx.export 接口，用起来不复杂。但有几个关键点，我建议你记牢。
- 加载训练好的 PyTorch 模型，并设置为
eval模式。 - 构造一个 dummy 输入（形状、数据类型要和真实推理时一致）。
- 调用
torch.onnx.export，指定输出路径、输入输出名称、动态轴等。 - 用
onnxruntime或onnx.checker验证导出的模型。
下面是一个典型的导出代码。嗯，这里要注意：dummy 输入的 batch size 最好用 1，动态 batch 后面再配。
import torch
import torchvision.models as models
# 1. 加载模型
model = models.resnet50(pretrained=True)
model.eval()
# 2. 构造 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)
# 3. 导出 ONNX
torch.onnx.export(
model,
dummy_input,
"resnet50.onnx",
input_names=["input"],
output_names=["output"],
dynamic_axes={
"input": {0: "batch_size"},
"output": {0: "batch_size"}
},
opset_version=11,
do_constant_folding=True
)
print("ONNX 导出成功！")do_constant_folding 建议打开，能减小模型体积。
3.3 ONNX 模型验证
导出 ONNX 后，别急着转 OM。先跑一遍验证，确保计算图没问题。我一般用 onnxruntime 做推理比对。
import onnx
import onnxruntime as ort
import numpy as np
# 检查模型结构
onnx_model = onnx.load("resnet50.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 模型结构检查通过")
# 用 onnxruntime 推理
session = ort.InferenceSession("resnet50.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
result = session.run([output_name], {input_name: test_input})
print("ONNX 推理输出 shape:", result[0].shape)如果这一步报错，多半是算子不兼容或者动态轴没配好。回头检查导出参数就行。
3.4 ONNX 转昇腾离线模型（OM）—— ATC 工具
ATC（Ascend Tensor Compiler）是昇腾平台的核心转换工具。说白了，它把 ONNX 模型编译成昇腾芯片能直接执行的离线模型（.om）。
我刚开始用 ATC 时，最头疼的是参数太多。其实常用的就那几个，我帮你整理了一张表。
| 参数 | 说明 | 示例 |
|---|---|---|
--model |
输入的 ONNX 模型路径 | --model=resnet50.onnx |
--framework |
框架类型，ONNX 填 5 | --framework=5 |
--output |
输出 OM 模型路径（不含后缀） | --output=resnet50 |
--input_shape |
指定输入 shape，支持动态 batch | --input_shape="input:1,3,224,224" |
--soc_version |
昇腾芯片型号，如 Ascend310P3 | --soc_version=Ascend310P3 |
--log |
日志级别，debug/info/warning/error | --log=info |
一个典型的转换命令长这样：
atc --model=resnet50.onnx \
--framework=5 \
--output=resnet50 \
--input_shape="input:1,3,224,224" \
--soc_version=Ascend310P3 \
--log=infonpu-smi info 先查清楚芯片型号。
3.5 动态 batch 与多输入处理
实际业务中，输入 batch 大小经常变化。ATC 支持动态 batch，但需要在转换时指定。
atc --model=resnet50.onnx \
--framework=5 \
--output=resnet50_dynamic \
--input_shape="input:-1,3,224,224" \
--dynamic_batch_size="1,2,4,8" \
--soc_version=Ascend310P3注意：--dynamic_batch_size 后面列出的 batch 值，必须是模型实际会用到的。我建议只列你需要的，别贪多，否则编译时间会变长，模型体积也会增大。
如果是多输入模型（比如双流网络），用 --input_shape 分别指定就行：
--input_shape="input1:1,3,224,224;input2:1,3,112,112"3.6 转换后的验证
拿到 .om 文件后，可以用 omg 工具或者直接写推理代码验证。我个人习惯用 MindSpore Lite 或 ACL 加载 OM 模型，跑一遍推理，对比输出误差。
1. 用 PyTorch 跑一次推理，保存输出。
2. 用 ATC 转换后的 OM 模型，在昇腾设备上跑同样的输入。
3. 计算两个输出的余弦相似度或最大绝对误差。
4. 误差在 1e-3 以内，基本就没问题。
我在项目中遇到过，转换后的模型输出和原始 PyTorch 结果差了几个百分点。后来发现是 ONNX 导出时 opset_version 太低，某些算子精度不对。升级到 opset=13 后，误差就消失了。所以验证这步，千万别省。