从 PyTorch 导出 ONNX

ONNX Runtime-MUSA 推理的是 ONNX 格式的模型。本章给出一条完整链路的示例：加载 PyTorch 模型 → 导出 ONNX → 用 ONNX Runtime-MUSA 推理 → 与 PyTorch 的输出比对，并说明导出时的常见注意事项。

以 ONNX Runtime 1.23.0 + ep 镜像为例，其已预装 PyTorch 2.7.1、torchvision 0.22.1、transformers 4.50.2，本章示例可在镜像内直接运行。容器的启动方式见启动容器，示例需要访问外网下载模型权重，docker run 时请加 --network host。

Step1 安装导出所需的 onnx 包

导出接口在写出模型文件时需要 onnx 包，镜像内未预装（推理本身不需要它）：

pip install onnx "numpy==1.23.5"

预期输出末行以 Successfully installed 开头，其中含 onnx，例如：

Successfully installed ml_dtypes-0.5.4 onnx-1.22.0 protobuf-7.36.2

注意：同时安装的其余包及其版本随时间变化，无需核对。装完确认 numpy 未被升级：

cd /tmp && python3 -c "
import numpy, onnx, torch, onnxruntime as ort
print('numpy :', numpy.__version__)
print('onnx  :', onnx.__version__)
print('torch :', torch.__version__)
print('ort   :', ort.__version__)
"

预期输出：

numpy : 1.23.5
onnx  : 1.22.0
torch : 2.7.1
ort   : 1.23.0+musa.26d77422

注意

必须同时指定 numpy==1.23.5。

直接执行 pip install onnx 会将 numpy 升级到 2.x，而镜像内的 PyTorch 与 ONNX Runtime 都是基于 numpy 1.23.5 构建的，升级后两者都将无法导入，报 A module that was compiled using NumPy 1.x cannot be run in NumPy 2.2.6。

环境无外网时，先在一台联网、架构与 Python 版本相同的机器上执行 pip download onnx -d wheelhouse/，将整个 wheelhouse/ 拷入后用 pip install --no-index --find-links wheelhouse/ onnx 安装。

Step2 导出一个单输入模型

下例使用 torchvision 的 ResNet-18：加载预训练模型、用 PyTorch 跑一次推理留作参照、导出 ONNX、再用 ONNX Runtime-MUSA 跑同一份输入并比对两者输出。

import torch
import torchvision
import numpy as np
import onnxruntime as ort

# 1. 加载模型，切换到推理模式
model = torchvision.models.resnet18(
    weights=torchvision.models.ResNet18_Weights.DEFAULT
).eval()

# 2. 准备一份示例输入，其形状即导出后模型的输入形状
x = torch.randn(1, 3, 224, 224)

# 3. 用 PyTorch 跑一次，作为比对参照
with torch.no_grad():
    torch_out = model(x)
print("torch output :", tuple(torch_out.shape))

# 4. 导出 ONNX
torch.onnx.export(
    model,
    (x,),
    "resnet18.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=17,
)
print("exported     : resnet18.onnx")

# 5. 用 ONNX Runtime-MUSA 跑同一份输入
session = ort.InferenceSession(
    "resnet18.onnx",
    providers=[("MUSAExecutionProvider", {"device_id": "0"}), "CPUExecutionProvider"],
)
onnx_out = session.run(None, {"input": x.numpy()})[0]
print("providers    :", session.get_providers())
print("onnx output  :", onnx_out.shape)

# 6. 比对两者输出
a = onnx_out.ravel().astype(np.float64)
b = torch_out.numpy().ravel().astype(np.float64)
cosine = float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))
print("cosine       : %.6f" % cosine)
print("argmax       : torch %d / onnx %d" % (int(b.argmax()), int(a.argmax())))

预期输出：

torch output : (1, 1000)
exported     : resnet18.onnx
providers    : ['MUSAExecutionProvider', 'CPUExecutionProvider']
onnx output  : (1, 1000)
cosine       : 0.999991
argmax       : torch <类别编号> / onnx <同一个类别编号>

注意：cosine 不低于 0.999、且两侧 argmax 相同，即表示导出与推理都正确。cosine 的具体数值随随机输入与运行环境小幅波动（实测在 0.99999 与 1.000000 之间），argmax 的取值随输入变化，只需两侧一致。首次运行会下载约 45 MB 的预训练权重。

Step3 导出参数说明

上例中 torch.onnx.export 的参数含义如下：

参数	说明
第 1 个参数	待导出的模型，导出前须调用 .eval()。训练模式下 Dropout、BatchNorm 的行为与推理不同，导出后结果会与预期不符
第 2 个参数	示例输入，写成元组。导出过程会用它实际跑一遍模型来记录计算图，因此它的形状与 dtype 决定了导出模型的输入规格
第 3 个参数	输出的 ONNX 文件路径
input_names / output_names	输入输出的名称。推理时按这些名称传入与读取，建议显式指定，否则会生成 onnx::Gemm_0 这类自动名称
opset_version	ONNX 算子集版本，建议 17。过低会导致新算子无法表达，过高则可能超出运行时支持范围
dynamic_axes	声明哪些维度在推理时可变，默认全部固定。见下一节

导出时用 torch.no_grad() 跑参照推理，可避免构建反向图、减少显存占用。

Step4 需要可变维度时：dynamic_axes

不声明 dynamic_axes 时，导出模型的输入形状与示例输入完全一致，换一个形状就会报错。若批量大小在业务中会变化，按下例声明：

import torch
import torchvision
import numpy as np
import onnxruntime as ort

model = torchvision.models.resnet18(
    weights=torchvision.models.ResNet18_Weights.DEFAULT
).eval()
x = torch.randn(4, 3, 224, 224)

torch.onnx.export(
    model,
    (x,),
    "resnet18_dynamic.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    opset_version=17,
)

session = ort.InferenceSession(
    "resnet18_dynamic.onnx",
    providers=[("MUSAExecutionProvider", {"device_id": "0"}), "CPUExecutionProvider"],
)
print("batch=4 :", session.run(None, {"input": x.numpy()})[0].shape)
print("batch=1 :", session.run(None, {"input": torch.randn(1, 3, 224, 224).numpy()})[0].shape)

预期输出：

batch=4 : (4, 1000)
batch=1 : (1, 1000)

注意：同一个 session 可以接受不同的批量大小。dynamic_axes 只声明确实需要变化的维度，其余维度保持固定——维度固定的模型便于引擎做形状相关的优化，通常也更快。

Step5 多输入模型示例：HuggingFace BERT

文本类模型通常有多个输入。下例从 HuggingFace 下载 BERT，把文本补齐到固定长度 128 后导出，并按名称传入三个输入：

import torch
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer, AutoModel

name = "google-bert/bert-base-chinese"
tokenizer = AutoTokenizer.from_pretrained(name)
model = AutoModel.from_pretrained(name).eval()

# 补齐到固定长度，导出模型的序列长度即为 128
enc = tokenizer("摩尔线程 GPU 推理", return_tensors="pt",
                padding="max_length", max_length=128, truncation=True)

with torch.no_grad():
    torch_out = model(**enc).last_hidden_state

torch.onnx.export(
    model,
    (enc["input_ids"], enc["attention_mask"], enc["token_type_ids"]),
    "bert.onnx",
    input_names=["input_ids", "attention_mask", "token_type_ids"],
    output_names=["last_hidden_state"],
    opset_version=17,
)

session = ort.InferenceSession(
    "bert.onnx",
    providers=[("MUSAExecutionProvider", {"device_id": "0"}), "CPUExecutionProvider"],
)
onnx_out = session.run(None, {k: v.numpy() for k, v in enc.items()})[0]

a = onnx_out.ravel().astype(np.float64)
b = torch_out.numpy().ravel().astype(np.float64)
print("onnx output :", onnx_out.shape)
print("cosine      : %.6f" % float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b))))

预期输出：

onnx output : (1, 128, 768)
cosine      : 0.999997

注意：cosine 不低于 0.999 即表示导出与推理正确，具体数值随运行环境小幅波动。导出多输入模型时，第 2 个参数按模型 forward 的形参顺序给出，input_names 与之一一对应；推理时按名称传入，顺序不限。

容器内访问 huggingface.co 受限时，启动容器加上环境变量 -e HF_ENDPOINT=https://hf-mirror.com，transformers 会改从镜像站下载。

注意事项

事项	说明
导出后必须验证精度	导出是一次模型转换，算子表达方式可能变化。上线前按模型集成的方法，用同一份输入对比 ONNX 与原模型的输出
优先使用固定形状	仅对业务上确实可变的维度声明 dynamic_axes，其余维度固定
导出前调用 .eval()	否则 Dropout、BatchNorm 仍按训练模式计算
权重超过 2 GB	导出时加 save_as_external_data=True，权重会写入模型文件旁的独立数据文件，两者需一并拷贝
前后处理不随模型导出	图像缩放、归一化、分词、检测框解码等仍在业务代码中完成，与推理解耦
导出在容器内完成	示例依赖镜像内的 PyTorch 与 ONNX Runtime 版本，在宿主机执行会因缺少相应环境而失败

导出得到的 ONNX 模型即可按模型集成接入自有工程。