英译中模型迁移ONNX：从PyTorch到onnxruntime的CPU推理优化实战

搞多语种内容平台这两年，后台最离不开的服务就是英译中。最开始我直接从 HuggingFace 拉模型，用 transformers 的 pipeline 几行代码就能跑，开发期确实舒服。但一上线问题就来了：线上机器不想装完整的 PyTorch 全家桶，镜像体积直线上涨；而且翻译服务主要跑在 CPU 上，PyTorch 原生推理既慢又浪费资源。于是我把 Helsinki-NLP/opus-mt-en-zh 这个英译中模型迁到了 ONNX 格式，用 onnxruntime 做推理，整体依赖和资源开销都降了一大截，量化之后体积和延迟更是进一步压缩。这篇文章把迁移过程中最关键的模型选型、导出参数、量化细节和生产落地经验完整写出来，希望能让你少走弯路。

1. 思路拆解：为什么要把翻译模型搬上 ONNX

1.1 线上服务怕的不是模型大，是依赖重

很多人以为模型上线最难的是模型本身，真正头疼的是依赖链。一个 transformers 加 torch 的 Python 服务，打包进 Docker 镜像通常在 2GB 以上，传到内网仓库、拉起容器、每次升级都要忍受这个体积。更别提 torch 的线程库和机器的 CPU 特性绑定很深，容器环境一变，推理性能波动就很大。ONNX 是一种开放模型格式，模型转出来后跟训练框架完全解耦，运行时只需要 onnxruntime 一个依赖，镜像体积能砍到原来的五分之一以下。

翻译场景还有个现实需求：服务可能会部署到多个区域、多规格的机器，甚至未来挪到边缘设备。ONNX 的跨平台能力很强，Windows、Linux、ARM 都有对应的 runtime，模型文件本身不绑定任何训练框架，这是 PyTorch 原生格式做不到的。后来我发现，不只是翻译模型，语音合成、视觉检测这类模型走 ONNX 路线的逻辑也基本一致：把训练产物沉淀成运行时无关的图格式，剩下的交给专用推理引擎去优化。

1.2 迁移方案的取舍

把 HuggingFace 模型迁移到 ONNX 有两条主流路线。一条是用 transformers.onnx 或 optimum 的 CLI 自动转换，省事但灵活性差；另一条是用 torch.onnx.export 手动导出 encoder 和 decoder，可控性强，适合要上生产、要定制量化的项目。我两条都走了：先用 optimum 快速验证可行性，再用 torch.onnx.export 重建，后面所有量化和服务化都是基于手动导出的这套文件做的。

这里要特别提醒：ONNX 导出和量化本身不难，难的是对模型结构的理解。MarianMT 是标准的 encoder-decoder 结构，翻译时 encoder 只对源句跑一次，把整句话编码成一组隐向量；decoder 每次只看一个目标 token，结合隐向量预测下一个 token，循环往复直到遇到结束符。如果你对这套生成流程没有概念，后面排查各种维度错误会非常痛苦。先把这一步想清楚，后面所有问题都好解。

2. 环境准备与模型选型

2.1 模型选型：为什么用 opus-mt-en-zh

HuggingFace 上英译中的模型不算少，我选了 Helsinki-NLP/opus-mt-en-zh。它是 MarianMT 架构，encoder 和 decoder 各 6 层，权重体积在几百 MB 量级，CPU 上推理压力不大，翻译质量对一般的内容场景完全够用。相比之下 NLLB-200 这类多语言模型虽然强，但动辄好几个 GB，对小服务来说不太划算。

选模型时建议先想清楚场景：是短句为主还是长文为主？对延迟敏感还是对质量敏感？纯 CPU 服务的话，模型体积直接决定成本。opus-mt-en-zh 还有个好处是 tokenizer 和模型行为都很常规，导出 ONNX 时基本不会遇到奇怪算子，对新手非常友好。如果你的业务对术语质量要求更高，也可以考虑在导入模型后用领域语料做一次微调，再走后面这套导出流程，逻辑是一样的。

2.2 依赖版本与 HuggingFace 下载

我的环境是 Python 3.10，装齐下面这套依赖：

版本别太激进，我试过 transformers 4.44 加 optimum 1.22 的组合，导出 MarianMT 时出现过兼容性问题，后来锁回上面这组，一切正常。环境装好后先下载模型：

export HF_ENDPOINT=https://hf-mirror.com huggingface-cli download Helsinki-NLP/opus-mt-en-zh --local-dir ./opus_mt_en_zh

直接从 HuggingFace 拉模型经常超时，把下载源切到国内镜像能明显提速，下载完的目录结构跟官方一致，后续 from_pretrained 直接指本地路径即可。这一步是整个流程里最不用动脑、但又最容易卡人的地方，先把它解决再谈导出。

3. 导出 ONNX 的实操过程

3.1 用 optimum-cli 快速导出

最快验证方案可行性的方式，是装好 optimum 后一行命令完成导出：

optimum-cli export onnx --model ./opus_mt_en_zh ./opus_onnx

导出完成后会得到 encoder_model.onnx、decoder_model.onnx、decoder_model_merged.onnx，以及 config.json 和 tokenizer 相关的配置、词表文件。其中 decoder_model_merged 是把 past_key_values 缓存逻辑合并进图的版本，推理时不用每次从头算 decoder 的注意力，长句翻译时速度提升非常明显。

如果只是验证方案，直接配合 optimum 的 ORTModelForSeq2SeqLM 类就能跑起来：

from optimum.onnxruntime import ORTModelForSeq2SeqLM model = ORTModelForSeq2SeqLM.from_pretrained("./opus_onnx")

之后可以像普通 transformers 模型一样接 pipeline。但注意，optimum 自动导出把很多配置固化在 config.json 里，后面对模型做 INT8 量化或自定义推理时，文件结构越复杂越容易出错。所以我后续生产用的手动导出方案，下面详细说。

3.2 手动 torch.onnx.export 兜底

手动导出前要先理解 MarianMT 的输入输出。encoder 输入是 input_ids 和 attention_mask，输出是 encoder_hidden_states；decoder 输入是 decoder_input_ids、encoder_hidden_states、encoder_attention_mask 和 decoder_attention_mask，输出是 logits。由于 ONNX 图不方便处理循环，我们只导出单步推理的 encoder 和 decoder，生成循环在外部用 Python 控制。

import torch from transformers import MarianMTModel model = MarianMTModel.from_pretrained("./opus_mt_en_zh") model.eval() class EncoderWrapper(torch.nn.Module): def __init__(self, m): super().__init__() self.enc = m.get_encoder() def forward(self, input_ids, attention_mask): return self.enc(input_ids=input_ids, attention_mask=attention_mask)[0] class DecoderWrapper(torch.nn.Module): def __init__(self, m): super().__init__() self.dec = m.get_decoder() def forward(self, decoder_input_ids, encoder_hidden_states, encoder_attention_mask, decoder_attention_mask): return self.dec( input_ids=decoder_input_ids, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, attention_mask=decoder_attention_mask, use_cache=False, )[0] enc_wrapper = EncoderWrapper(model) dec_wrapper = DecoderWrapper(model) dummy_ids = torch.tensor([[model.config.decoder_start_token_id]], dtype=torch.int64) dummy_enc = torch.randn(1, 32, model.config.d_model) dummy_attn = torch.ones(1, 32, dtype=torch.int64) dummy_dec_attn = torch.ones_like(dummy_ids) torch.onnx.export( enc_wrapper, (torch.ones(1, 32, dtype=torch.int64), dummy_attn), "encoder.onnx", input_names=["input_ids", "attention_mask"], output_names=["encoder_hidden_states"], dynamic_axes={ "input_ids": {0: "batch", 1: "seq"}, "attention_mask": {0: "batch", 1: "seq"}, }, opset_version=14, ) torch.onnx.export( dec_wrapper, (dummy_ids, dummy_enc, dummy_attn, dummy_dec_attn), "decoder.onnx", input_names=["decoder_input_ids", "encoder_hidden_states", "encoder_attention_mask", "decoder_attention_mask"], output_names=["logits"], dynamic_axes={ "decoder_input_ids": {0: "batch", 1: "dec_seq"}, "encoder_hidden_states": {0: "batch", 1: "enc_seq"}, "encoder_attention_mask": {0: "batch", 1: "enc_seq"}, "decoder_attention_mask": {0: "batch", 1: "dec_seq"}, }, opset_version=14, )

这里有个容易踩坑的点：MarianDecoder 的第一个 attention_mask 参数指的是 decoder 自身的自注意力掩码，encoder 的掩码要单独放在 encoder_attention_mask 参数里。我第一次导出时把两者搞反了，结果推理出来的翻译完全是乱码，后来排查才发现是语义反了。导出完可以用 onnx.checker.check_model 校验一遍，避免低级错误。

3.3 导出参数背后的道理

dynamic_axes 是 ONNX 导出中最重要的配置。翻译场景里每次输入长度都不一样，decoder 循环里序列还会越变越长，所以必须把 batch 和 seq 两个维度都标成动态；如果只标 seq，后面想批处理时就会报维度不匹配。

opset 版本我选 14。ONNX 对算子版本的要求一直在涨，opset 太老会触发算子 fallback，运行时效率低；opset 太新则要求 onnxruntime 版本足够新，一些老机器上的 runtime 可能不支持。14 是目前兼容性较好的折中档，我在这套环境上没遇到任何算子兼容问题。do_constant_folding 默认开启，可以折叠一部分常量算子，对推理有正面帮助，不需要额外处理。

另外建议导出前设置 torch.set_num_threads(1) 或 OMP_NUM_THREADS 环境变量，让 trace 过程保持单线程。否则 torch 会自动开满 CPU 线程，导出时间反而变长，还容易出现莫名其妙的卡顿。

4. ONNX 推理与结果对齐

4.1 onnxruntime 下的推理循环

ONNX 图本身不会替你管理生成循环，翻译时要在外面手动模拟 generate。核心流程是：先用 encoder 把源句编码成隐向量，再拿 decoder 一步步产出目标 token，直到遇到 eos 或达到最大长度。

import numpy as np import onnxruntime as ort from transformers import MarianTokenizer tokenizer = MarianTokenizer.from_pretrained("./opus_mt_en_zh") enc_sess = ort.InferenceSession("encoder.onnx", providers=["CPUExecutionProvider"]) dec_sess = ort.InferenceSession("decoder.onnx", providers=["CPUExecutionProvider"]) def translate(text, max_len=64): tok = tokenizer(text, return_tensors="np", max_length=512, truncation=True) enc_out = enc_sess.run(None, { "input_ids": tok["input_ids"], "attention_mask": tok["attention_mask"], })[0] dec_ids = np.array([[model.config.decoder_start_token_id]], dtype=np.int64) enc_attn = tok["attention_mask"] for _ in range(max_len): dec_attn = np.ones_like(dec_ids) logits = dec_sess.run(None, { "decoder_input_ids": dec_ids, "encoder_hidden_states": enc_out, "encoder_attention_mask": enc_attn, "decoder_attention_mask": dec_attn, })[0] next_id = np.argmax(logits[:, -1, :], axis=-1, keepdims=True) if next_id.item() == model.config.eos_token_id: break dec_ids = np.concatenate([dec_ids, next_id], axis=1) return tokenizer.decode(dec_ids[0], skip_special_tokens=True)

这段代码对应的是手动导出的 encoder.onnx 和 decoder.onnx，不依赖 optimum 运行时。如果你用的是 optimum 导出的 decoder_model_merged.onnx，输入会多出一长串 past_key_values 相关张量，纯 ORT 写起来会繁琐不少。先跑通这段基础版，再考虑加缓存优化，这个顺序能省大量调试时间。session 创建一次放在模块级，不要在每次请求里重复初始化，这点在后面的部署部分还会再强调。

4.2 与 PyTorch 的翻译结果对比

模型迁移最怕结果对不上。我把同一批句子分别用 transformers 的 generate 和上面这段 ONNX 推理跑了一遍，对比了 20 个句子，绝大多数字面完全一致，个别句子因为浮点精度和算子合并出现了同义词级别的差异，比如 “the” 翻成 “一个” 而不是 “这个”，语义完全不受影响。

验证时建议先固定 greedy search（num_beams=1），别急着上 beam search。beam search 要在外部自己维护多条候选序列，复杂度高不少；先保证 greedy 结果一致，能更精准判断问题出在导出环节还是量化环节。等确认推理正确后，需要更高的翻译质量再回来做 beam search，也是合理的迭代路径。

5. INT8 量化：体积与延迟一起压

5.1 动态量化的实操

ONNX 模型量化有两条路线：动态量化和静态量化。翻译模型这种输出张量粒度比较细的场景，静态量化需要准备校准数据集，对数据分布敏感，调起来费劲；我实际生产中用动态量化，只把权重转成 INT8，激活值在推理时动态计算。效果上体积和延迟都有明显收益，对精度影响却很小。

from onnxruntime.quantization import quantize_dynamic, QuantType for src, dst in [ ("encoder.onnx", "encoder_q8.onnx"), ("decoder.onnx", "decoder_q8.onnx"), ]: quantize_dynamic(src, dst, weight_type=QuantType.QInt8)

QuantType 我选了 QInt8。QUInt8 在部分老 CPU 上可能更快，但兼容性和数值表现不如 QInt8 可控，我实测下来 QInt8 更放心。量化参数里还有个 per_channel，默认 False，按整个权重张量算一个缩放系数；如果量化后精度损失明显，把它改成 per_channel=True，量化粒度更细，精度损失更小，缺点是文件略大。我这边开启 per_channel 后翻译结果和 FP32 几乎一致。

5.2 量化前后的数据对比

我当时的运行环境是 8 vCPU、内存上限 4GB 的容器，翻译“The quick brown fox jumps over the lazy dog.”，最大生成长度 32，预热后连续测 20 次取平均：

数字只是相对参考，不同机器差异很大。关键信息是：体积降到四分之一左右，延迟大约减少四成，内存峰值接近腰斩。对纯 CPU 小容器来说，这个差距有时候就是一台机器钱的区别。精度方面我拿 200 条新闻标题做了对比，INT8 和 FP32 只有不到 5% 的句子存在用词差异，没有发现翻错语义的情况，完全够用。

6. 生产部署实战与常见坑

6.1 服务化落地：session 复用与线程模型

模型文件准备好后，我用 FastAPI 包了一层 HTTP 服务。这里最重要的不是接口代码，而是 onnxruntime session 的使用方式：全局只初始化一次，多个请求线程共享同一个 session。onnxruntime 的 session 是线程安全的，重复创建 session 反而会反复加载模型、浪费几百毫秒。我最初实现时把 session 创建写进了请求函数里，压测一上来并发，延迟立刻飙升，改成模块级懒加载后问题消失。

from functools import lru_cache import onnxruntime as ort @lru_cache(maxsize=1) def get_sessions(): enc = ort.InferenceSession("encoder_q8.onnx", providers=["CPUExecutionProvider"]) dec = ort.InferenceSession("decoder_q8.onnx", providers=["CPUExecutionProvider"]) return enc, dec

还要注意容器启动后主动做一次暖机推理。我踩过坑：ONNX Runtime 第一次跑某个 shape 的图时，初始化开销不小，如果正好赶上请求高峰，线上就会出现明显毛刺。启动时拿一条固定句子先跑一次，把初始化成本提前消耗掉，这个动作效果很直接。另外建议所有 session 显式指定 providers，避免不同环境的默认 provider 不一致导致性能差异。

6.2 常见问题速查表

这些坑我基本都真实踩过。尤其是第一个，新手最容易卡住：导出是单线程 trace，不需要 torch 开一堆线程，加上环境变量限制就好。顺带一提，这套用 onnxruntime 服务模型的思路并不局限于翻译，语音合成比如 sherpa-onnx TTS、视觉检测比如 YOLO 导出 ONNX，走的也是同一套路线：先拿到标准格式的图，再用轻量运行时接管推理，最后按需量化部署。核心方法是一致的。

把翻译模型迁到 ONNX 这件事，我前后折腾了差不多两个星期。最开始以为难点在导出命令上，后来才发现真正花时间的是理解 encoder-decoder 推理流程、设计动态维度、以及处理量化和部署里的各种细节。现在这套方案在线上已经跑了一段时间，响应速度和服务资源占用都符合预期。如果你也在做类似的事，我的建议是：先用 optimum-cli 快速跑通全流程建立信心，再手动导出做精细控制；量化优先选动态量化；部署前一定做好 session 复用和暖机。把这些点做到位，ONNX 迁移这条路的性价比确实很高。