最近在搞一个批量翻译工具需求很朴素把一堆英文技术文档批量翻成中文。团队里本来就习惯用 HuggingFace 生态我自然先从上面找现成的英译中模型跑通验证之后再做部署。最初直接用 PyTorch 权重推理翻译质量是没得说可真到了要上线的时候问题接踵而来——生产机上环境依赖重、模型加载慢、显存占用也不低。折腾了几天我决定把 HuggingFace 模型迁移到 ONNX底层推理换成 ONNX Runtime 来跑。这篇文章就把从选型、导出、量化到部署的完整过程记录下来踩过的坑也一并列出来给正在做类似迁移的朋友一个参考。1. 为什么要折腾 ONNX 迁移1.1 直接死在 PyTorch 推理上的教训先说结论不是 PyTorch 不好而是生产环境对依赖体积、启动速度和推理性能的要求比实验室苛刻得多。我当时用 transformers 库加载一个开源团队发布的英译中翻译模型推理链路很简单tokenizer 编码 → model.generate() → tokenizer 解码。跑通 Demo 很容易但上线前盘点资源时发现几个硬伤。首先是依赖体积。单是 PyTorch 的 CUDA 版本安装包就超过 2GB再加上 transformers、tokenizers、sentencepiece 等一连串依赖整个 Python 环境轻松逼近 5GB。对这个要交付的 Docker 镜像来说体积是非常不友好的。CI 构建时间、镜像仓库存储、服务器磁盘占用全部被放大。其次是启动速度。transformers 加载模型要做 lazy initialization 和权重映射实测冷启动加载一个 300M 左右的翻译模型在普通 CPU 机器上要 20 秒以上GPU 机器也得 10 秒左右。如果服务要频繁扩容缩容、快速拉起新实例这部分时间完全是浪费。最关键的是推理性能。transformers 的 generate 接口虽然方便但内部有大量灵活性开销。同样的模型权重迁移到 ONNX 导出后在 CPU 上推理速度提升了大约 1.5 到 2 倍在 GPU 上配合 CUDA EP 和半精度提升更明显。原因在于 ONNX Runtime 会做算子融合、内存复用和常量折叠这些优化是 PyTorch eager 模式很难做到的。如果你只是为了本地跑几个 Demo、验证模型效果PyTorch 完全够用。但一旦要多实例部署、要控制成本、要降低响应延迟ONNX 迁移带来的收益非常直观。1.2 迁移到 ONNX 之后到底拿到了什么ONNX 是一个开放模型中间表示标准核心价值在于“一次导出到处推理”。模型保存成 ONNX 格式后可以脱离原始训练框架运行任何支持 ONNX 的推理引擎都能加载。对我这个翻译项目来说最关心的几点依赖大幅缩减。推理阶段只需要 onnxruntime 运行时不需要 torch也不需要把 transformers 全家桶带进生产环境。CPU 推理性能更好。ONNX Runtime 对算子做了大量融合优化尤其适合 CPU 环境下的文本生成。支持量化。int8 量化后模型体积可以缩小到原来的四分之一左右推理延迟还能进一步降低这对低配服务器非常香。跨平台部署。同一个 ONNX 文件可以跑在 Linux、Windows、macOS甚至 Android/iOS 上一套导出多处使用。用一句话类比PyTorch 模型像一辆改装赛车性能上限高但只能在特定场地跑ONNX 像标准集装箱虽然不能随意改装但全世界都有配套的港口和运输系统。但这不代表迁移是零成本的。文本生成类模型有自己的特殊性导出时要处理动态长度、自回归循环、beam search 等逻辑不能像图像分类模型那样一导了之。这也是这篇文章存在的意义。2. 选模型、备环境、加载权重2.1 英译中模型怎么选HuggingFace 上的翻译模型不少我在这个项目里实际评估过三类先放个对比表模型参数量架构优点缺点Helsinki-NLP/opus-mt-en-zh约 300MMarianMT英中专精、体积小、推理快长句翻译偶尔有漏译facebook/m2m100-418M418MEncoder-Decoder多语言互译、效果均衡对英中任务来说资源开销偏大facebook/nllb-200-distilled-600M600MEncoder-Decoder支持超多语言模型重导出算子和部署成本高最终我选了 Helsinki-NLP 团队的 opus-mt-en-zh。理由很简单我的场景就是英译中不需要多语言能力用一个小而专的模型导出、部署、运维成本都最低。它的底层架构是 MarianMT本质上是标准的 Encoder-Decoder TransformerONNX 导出路径相对成熟网上能查到的资料也多一些。挑模型时还有两个建议。第一不要只看翻译质量分数还要看权重文件大小、分词器依赖是否复杂。有些模型分数高但依赖特殊的前处理逻辑导出 ONNX 时很容易遇到算子不支持的坑。第二优先选 transformers 官方支持较好的模型类型比如 MarianMT、M2M100、NLLB 这些因为 transformers 的自动导出工具对它们适配更好遇到问题也更容易搜到解决方案。2.2 环境与依赖配置我用的环境是 Python 3.10 Ubuntu 20.04GPU 机器是 CUDA 11.8CPU 机器是 8 核的普通云主机。先列一下依赖pip install torch2.1.0 transformers4.38.0 pip install onnx1.15.0 onnxruntime1.17.0 pip install onnxruntime-gpu1.17.0 # GPU 机器选装 pip install optimum[onnxruntime] # 用于 optimum-cli 一键导出版本这里要特别注意transformers 版本和 torch 版本会影响导出的计算图细节建议尽量用新一点的版本。我最早用 transformers 4.30 导出时注意力算子在 ONNX 图的输出 shape 处理上有问题导致推理结果全错排查了很久才发现是库版本太旧。升级 transformers 之后同样的代码就导出了正确结果。GPU 机器装 onnxruntime-gpu 而不是 onnxruntime这两个包不能同时存在否则运行时容易挂错后端。如果只是 CPU 部署只装 onnxruntime 就够了体积小很多。2.3 把模型和分词器顺利加载起来加载模型很简单但有几个细节不能省from transformers import AutoTokenizer, AutoModelForSeq2SeqLM model_name Helsinki-NLP/opus-mt-en-zh tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForSeq2SeqLM.from_pretrained(model_name) model.eval()第一一定要调用model.eval()。如果忘记切到 eval 模式dropout 层仍然处于训练状态导出的 ONNX 图会包含随机行为推理结果不稳定而且这个问题非常隐蔽导出的模型不会报错只是结果偶发异常。第二如果网络条件不太稳定、权重文件比较大建议先手动把模型文件下载到本地目录再从本地路径加载。比如把文件放在./models/opus-mt-en-zh/目录下加载时用AutoModelForSeq2SeqLM.from_pretrained(./models/opus-mt-en-zh)既能避免反复拉取远端文件也方便后续离线部署。第三加载完模型后先做一次最小推理确认模型和分词器本身没有问题再进入导出环节。这一步能过滤掉一大半“模型没问题是导出后出错”的误判。3. 核心迁移从 PyTorch 导出 ONNX 的完整过程3.1 快速上手optimum-cli 一行命令导出如果你的目的是快速验证“这个模型能不能导出”我建议直接用optimum-cli它是 HuggingFace 官方推荐的模型导出工具对 transformers 模型的支持最完整。安装好optimum[onnxruntime]之后执行optimum-cli export onnx --model Helsinki-NLP/opus-mt-en-zh onnx/执行完成后onnx/目录下会生成模型文件和配置文件。这个工具会自动选择默认的输入输出包括 encoder 和 decoder 的单步前向还会处理分词器的特殊 token 映射。但说实话自动导出对文本生成模型只能导出“单次前向”也就是 encoder 和 decoder 的一步推理。完整的自回归生成循环包括逐 token 迭代、结束条件判断、beam search 等仍然需要自己在推理侧写代码。所以这个方法适合做第一时间验证真正要落地部署我还是推荐手动导出。3.2 手动导出翻译模型真正适配 ONNX 的正确姿势为什么要手动导出因为翻译模型生成时是自回归的每生成一个 token就要把新的 token 拼到 decoder 输入里再跑一次前向计算。transformers 把这一整套过程封装在model.generate()里但 ONNX 导出的是单步计算图我们只能导出“一次前向”的逻辑循环必须留在外部。手动导出的核心是把整个翻译模型包装成“单步前向”结构import torch from transformers import AutoTokenizer, AutoModelForSeq2SeqLM model_name Helsinki-NLP/opus-mt-en-zh tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForSeq2SeqLM.from_pretrained(model_name) model.eval() # 构造固定序列长度的示例输入 src_text This is a test sentence for machine translation. enc tokenizer(src_text, return_tensorspt, paddingmax_length, max_length128) input_ids enc[input_ids] attention_mask enc[attention_mask] # decoder 的起始 token 从模型 config 里取不要手写死 decoder_start_id model.config.decoder_start_token_id decoder_input_ids torch.tensor([[decoder_start_id]], dtypetorch.long) class TranslationWrapper(torch.nn.Module): def __init__(self, model): super().__init__() self.model model def forward(self, input_ids, attention_mask, decoder_input_ids): encoder_outputs self.model.get_encoder()(input_ids, attention_mask) decoder_outputs self.model.get_decoder()( decoder_input_ids, encoder_outputsencoder_outputs ) logits self.model.lm_head(decoder_outputs[0]) return logits wrapper TranslationWrapper(model) torch.onnx.export( wrapper, (input_ids, attention_mask, decoder_input_ids), translation_step.onnx, input_names[input_ids, attention_mask, decoder_input_ids], output_names[logits], dynamic_axes{ input_ids: {0: batch, 1: src_seq}, attention_mask: {0: batch, 1: src_seq}, decoder_input_ids: {0: batch, 1: dec_seq}, logits: {0: batch, 1: dec_seq}, }, opset_version14, )这段代码的核心就三件事把 encoder 和 decoder 串起来、把 lm_head 的输出暴露出来、设置动态轴。导出后的模型接收三个输入返回 logits这个 logits 形状是[batch, dec_seq, vocab_size]外层拿到 logits 后自己选下一个 token 就行。这里有个经验要重点说不要尝试把 beam search 或自回归循环塞进 ONNX 图里。虽然技术上可行但图的复杂度和调试难度会成倍增加收益却很有限。更合理的做法是让 ONNX 模型只负责“算 logits”循环逻辑留在部署层的 Python 或 C 代码里这几乎是我看到的所有生产项目的通用方案。3.3 动态轴配置才是关键动态轴dynamic_axes是文本模型导出的核心难点。如果你把 seq 维度固定死导出的模型只能翻译固定长度的句子这在真实场景中完全不可用。设置动态轴相当于告诉 ONNX这些维度在推理时是可变的具体数值由实际输入决定。比如input_ids: {0: batch, 1: src_seq}意思是输入的第一个维度是 batch size第二个维度是序列长度两个维度都可以在运行时指定。decoder 侧的dec_seq同理它会在生成过程中不断变长。但注意ONNX Runtime 对动态轴的支持不是无限制的。seq 长度决定了注意力矩阵的大小如果计算图中间某个算子只支持静态形状推理时就会报 shape mismatch 错误。我遇到这种情况时的排查思路是先把所有轴固定成静态形状确认模型功能正常再逐步打开动态轴定位是哪个算子不支持。另外如果你想把 KV Cache键值缓存也导出进来通常会把它保持固定形状因为 KV Cache 的扩容逻辑不适合用动态轴描述。这也是为什么“导出单步前向模型 外层自回归循环”成为主流做法。3.4 导出报错排查导出过程中我先后遇到过几类典型报错按频率从高到低列一下Unsupported operator。某个自定义算子或者较新的 PyTorch 操作没有对应的 ONNX 实现。解决办法通常是换用更基础的操作组合或者避免在导出路径中使用某些高级 API。遇到具体算子报错时去 ONNX 算子集文档里查一下是否支持比对着报错信息瞎猜高效得多。Shape inference 失败。dynamic_axes 与模型内部某个静态 shape 存在冲突。解决办法是把冲突的那个维度改成固定值或者缩小动态轴范围。比如有些模型内部把 seq 维度参与 reshape 时写死了这时候动态轴就不能覆盖那个维度。半精度权重导出异常。如果模型本身是 fp16 的建议先用 fp32 导出再在推理侧做量化不要在导出环节混合精度否则容易导出出精度异常的图。4. 优化int8 量化与推理性能调优4.1 量化类型怎么选导出的 ONNX 模型默认是 fp32体积大约和 PyTorch 权重差不多。但在 CPU 上做大规模部署时fp32 的算力开销和内存带宽都是瓶颈。量化就是把权重从 fp32 压缩到 int8用更少的位数存储和计算。以这个英译中模型为例fp32 权重文件接近 600MBint8 量化后大约 150MBCPU 推理延迟往往能再降 30%-50%。这里先澄清两个概念动态量化Dynamic Quantization是在推理时才把激活值转为 int8权重提前量化为 int8适合 NLP 模型静态量化Static Quantization需要提前准备校准数据集在导出前就确定激活值的缩放系数更适合 CV 模型。对翻译模型这种激活分布不太规律的逐 token 生成任务动态量化是更稳妥的起点实现也最简单。4.2 动态量化实操与效果验证ONNX Runtime 里做动态量化非常简洁from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_inputtranslation_step.onnx, model_outputtranslation_step_int8.onnx, weight_typeQuantType.QInt8, )跑完不到一分钟。但我强烈建议在量化前先跑通原始 fp32 模型的推理确认输入输出没问题否则量化模型出了 bug 很难判断是量化引入的还是模型本身有问题。量化后的模型要重新做质量验证。我当时用一组 100 句的测试集对比 fp32 和 int8 的输出大多数句子翻译结果几乎一致偶尔有几句话的选词略有差异但语义都能保持。如果量化后出现严重质量下降可以考虑保留部分敏感层为 fp32或者改用混合精度量化策略而不是直接放弃量化方案。4.3 图优化与内存占用优化除了量化ONNX Runtime 还支持图优化级别设置。默认情况下会开启全部优化包括算子融合、布局优化、冗余节点清理等。CPU 环境下这些优化对翻译模型的加速效果明显而且完全免费不需要额外配置。ONNX Runtime 的内存优化也很关键。推理时尽量复用输出 buffer避免每轮循环都重新分配内存。在自回归生成场景里decoder 输入长度逐轮增加如果每轮都新分配一份内存内存碎片和分配开销会拖慢整体速度。我的做法是在循环外预先分配最大长度所需的空间循环内只更新当前有效的部分。5. 部署实战用 ONNX Runtime 跑起翻译服务5.1 完整的自回归推理代码部署侧我用 Python onnxruntime。整个过程分两层外层负责自回归循环内层调用 ONNX 模型做单步计算。import onnxruntime as ort import numpy as np sess ort.InferenceSession(translation_step_int8.onnx) def translate(text: str, max_length: int 128): enc tokenizer(text, return_tensorsnp, paddingmax_length, max_length128) input_ids enc[input_ids].astype(np.int64) attention_mask enc[attention_mask].astype(np.int64) decoder_ids np.array([[model.config.decoder_start_token_id]], dtypenp.int64) for _ in range(max_length): inputs { input_ids: input_ids, attention_mask: attention_mask, decoder_input_ids: decoder_ids, } logits sess.run(None, inputs)[0] # [batch, dec_seq, vocab] next_token logits[:, -1, :].argmax(axis-1) # greedy decoder_ids np.concatenate([decoder_ids, next_token.reshape(-1, 1)], axis1) if next_token.item() tokenizer.eos_token_id: break return tokenizer.decode(decoder_ids[0], skip_special_tokensTrue)这是最简单的 greedy 解码。如果你需要 beam search要在每次迭代时维护多个候选序列把 logits 转成 log 概率再按 beam width 做剪枝。ONNX 模型不关心这些算法细节它只负责给你 logits周围全部是部署层的代码逻辑。5.2 性能对比与线上配置建议我拿 1000 句英文长度在 10-50 词的测试集在 8 核 CPU 机器上跑了一轮得到的典型数据如下推理方式平均吞吐句/分钟峰值内存MBPyTorch CPU1081200ONNX fp32 CPU204900ONNX int8 CPU348650需要说明具体数字会因模型和机器配置不同而浮动但趋势是稳定的——ONNX fp32 比 PyTorch CPU 快 1.5-2 倍int8 量化之后又提升 50% 以上。如果换 GPU 环境使用 onnxruntime-gpu CUDA EP配合半精度还能再快不少。线上部署的时候我给几个实用建议用 FastAPI 包一层 HTTP 接口把 ONNX 推理放在后台线程池里避免阻塞请求。模型在进程启动时加载一次后续请求复用同一个 InferenceSession不要每个请求都重新加载。对批量文本做长度分桶相似长度的句子 pad 到相近长度减少无效 padding 的计算浪费。进程内设置合适的OMP_NUM_THREADS不是线程数越多越快要和 CPU 核数匹配。5.3 从 Python 服务到更多部署形态ONNX 模型的一个大优势是部署形态非常多样。我自己的项目用的是 Python 服务但同一个 ONNX 文件也能用 C API、C# API 甚至移动端框架加载。ONNX Runtime 官方提供了多语言绑定推理性能和 Python 版本几乎一致。如果你面对的是嵌入式场景或者低资源环境思路其实和我做翻译模型迁移是一样的把一个训练好的框架模型导成 ONNX再针对目标平台做量化和图优化。比如语音合成和语音识别领域常见的 Sherpa 系列工具不少模型就是直接用 ONNX 格式发布的它们走的流程也类似。学会一套“框架模型 → ONNX → 优化 → 部署”的通用方法论以后遇到其他模型迁移就不会慌。6. 常见问题与排查实录6.1 输出结果全错先检查预处理一致性这类问题 90% 出在输入输出预处理不一致。检查三件事分词器是否和导出时用的是同一个版本padding 策略是否一致导出时是否做了model.eval()。我踩过一次很隐蔽的坑导出时用paddingTrue让 tokenizer 自动补到最长序列推理时忘了加 padding导致 seq 长度对不上。ONNX 模型虽然能跑但 padding mask 的位置密集失真输出的部分 token 已经完全不对了而且不会报任何 error排查难度超高。6.2 输入 dtype 报错int64 的坑ONNX Runtime 对 dtype 非常敏感。导出时输入默认是 int64推理时如果用 numpy 不小心传成 int32运行时会直接报 “Input ... is of type int32” 之类的错误。解决办法很简单构造输入时显式.astype(np.int64)。这一点我在代码示例里已经标出来了但每次写新脚本都容易忘建议把输入构造逻辑封装成单独函数。6.3 decoder 起始 token 错误翻译模型的初始 decoder 输入不是随便取的s或者eos每种模型的配置可能不同。我第一次导出时没细看 tokenizer直接用了eos_token_id初始化 decoder结果大量输出提前终止翻译出来的句子缺头少尾。后来改成从model.config.decoder_start_token_id读取问题立刻消失。这个教训就是初始化 decoder 的 token 一定要从 config 里拿不要猜、不要写死。6.4 beam search 与精度问题beam search 的可变状态维护是容易出错的地方。我的建议是分步调试先实现 greedy 解码并跑通再加 beam search。如果 beam search 结果异常优先检查工具栏logits 是否转换成了 log 概率是否加了长度惩罚不同 beam 之间的 mask 是否正确。算法逻辑建议先用纯 PyTorch 跑一遍验证确认没问题之后再替换底层为 ONNX Runtime这样出问题时更容易定位是算法问题还是模型问题。6.5 量化后性能提升不明显量化后性能没提升通常是两个原因。第一计算图里大量不可量化算子比如 LayerNorm 和 Softmax 在动态量化中默认保留 fp32如果这类算子在关键路径上占比较高整体收益自然有限。第二CPU 指令集不支持或未启用加速指令比如 VNNI/AVX512。同样是 int8 推理支持 VNNI 的 CPU 和不支持的 CPU性能差距可能在一倍以上。判断是不是这个原因可以用lscpu看标志位或者在 ONNX Runtime 日志里打开算子内核执行信息确认。整个流程走完我最深的一点体会是翻译模型迁移 ONNX最花时间的不是导出命令本身而是理解自回归模型的结构、动态轴怎么配、单步推理怎么设计。一旦把第一个模型完整跑通之后换成其他语种、其他模型基本就是改改模型名和分词器的事。希望这篇文章能帮你少走几趟弯路直接跳过那些只有踩过坑才会知道的门槛。