上个月我把HuggingFace上一个开箱即用的英译中模型迁到了ONNX,整个流程走下来,发现网上讲这个操作的文章大多是“给你一条命令,跑完收工”的流水账,真正踩坑的地方——比如动态序列长度、decoder缓存、量化后输出乱码——全都要靠自己去试错。这篇博文把我自己的完整过程、参数选择、问题排查都写下来,希望能帮到正在做类似迁移,尤其是第一次碰ONNX的人。
先说清楚这个活儿是什么:HuggingFace上有大量预训练好的翻译模型,我选的是Helsinki-NLP/opus-mt-en-zh,直接下载就能做人名招牌级的英译中,无需自己训练。但它默认跑在PyTorch上,生产环境里要么用TorchServe这类重型服务,要么得扛着一整套Python运行时部署,延迟和内存都不太友好。而模型迁移到ONNX之后,可以直接用ONNX Runtime推理,同样的CPU机器上延迟能压掉一半左右,而且部署文件只有几百MB里最小的那个,不需要再装PyTorch。
这篇文章的目标读者是有一定NLP基础、想深入了解模型落地的同学。如果你是第一次听说ONNX,前面的原理部分要细看;如果你已经被各种cast、shape mismatch折腾得快自闭了,直接跳到第五节查问题清单也行。
1. 迁移前的整体思路:为什么要从PyTorch走到ONNX
1.1 选型玄学:为什么我选Helsinki-NLP的opus-mt-en-zh
HuggingFace上英译中模型看着多,真正能拿来直接用的其实就那么几类。我推荐Helsinki-NLP/opus-mt-en-zh有两个原因。
第一,体积和速度的平衡很好。它是一个基于MarianMT架构的模型,参数量在300M以内(实际76M左右,具体看lib版本),在各种低配服务器上都能跑起来,不像NLLB-200那种动辄600M、甚至几个G的模型,虽然翻译质量天花板更高,部署成本也高了一个量级。第二,它是标准的Encoder-Decoder结构,输入不是那种只有unidirectional encoder的LM,输出也不是简单的“最后一个token接分类头”,这种结构在迁移ONNX时更典型——你需要同时导出encoder和decoder,还要处理past key values缓存,踩一遍坑之后,其他Seq2Seq模型迁移就有经验了。
再解释一下MarianMT这个架构,方便你之后看模型文件不懵。它和市面上大多数机器翻译模型一样,左边Encoder负责把英文句子编码成一个语义向量序列,右边的Decoder负责逐个生成中文token,每一步生成的token都会加入下一轮的输入。关键点在于,Decoder在推理时会复用前面所有时刻计算的Key和Value,也就是past key values缓存,否则每一步都要重新算前面所有词,复杂度从O(n)变O(n^2)。这个机制是后面ONNX图结构里最需要理解的部分。
1.2 ONNX Runtime到底比PyTorch好在哪里
很多人会有疑问:PyTorch也能在CPU上跑,为什么非要迁移到ONNX?我实测下来的核心差异有四个:
- 图优化与算子融合:ONNX Runtime在加载模型时会做算子融合、常量折叠、内存规划优化。PyTorch的eager execution是一步一步调内核,很多可以合并的操作被拆散了,比如LayerNorm里的多个小算子可能在ONNX里被融合成一个,内存读写开销明显下降。
- 没有Python解释器开销:PyTorch推理时Python代码和C++算子交替执行,Python解释器会卡在GIL上,多线程推理时会互相让锁。ONNX Runtime整个执行是C++完成的,Python只负责调用,服务端承载更高并发时不至于是解释器在拖后腿。
- 部署环境不再被Python绑架:你只需要一个opset兼容、带ONNX Runtime的环境,甚至连Python都可以不要,C++、Java、C#、Rust接口都全。这一点在嵌入式、桌面端、游戏引擎里做翻译功能时尤其香。
- 量化支持更成熟:PyTorch也有量化工具,但要走量化感知训练、动态量化、静态量化三条路,所需的数据管线更复杂。ONNX Runtime那边一个quantize_dynamic就能把fp32模型压成int8,速度和体积同时优化,后面会细讲。
当然,ONNX也不是银弹。你的模型如果结构里有非常新的自定义算子,opsets兼容可能跟不上;遇到torch.onnx.export失败时,要补onnxscript去做算子映射,这个我要在问题排查部分专门说。
1.3 迁移方案整体拆成四步
整个迁移不是一条optimum-cli命令就完事的。我自己的节奏是这样:
- 导出:把HuggingFace模型转换为ONNX格式的encoder和decoder文件,同时保存tokenizer和配置。
- 验证:用同一句输入分别跑PyTorch和ONNX,逐token对比输出概率,差异必须在可接受范围内。
- 优化:如果导出结果没问题,再做INT8量化和runtime配置优化。
- 部署:封装成推理接口,接入API或脚本,这个环节重点考虑的是缓存管理和beam search实现。
后面几节基本就是围绕这四步来展开,每一步搞不定都会在后头反噬。
2. 导出模型:Optimum CLI还是手写torch.onnx.export
2.1 工具选型:三选一,别偷懒
导出框架有好几种路径,我把它们的差异列一下,方便你按自己情况选:
| 工具路径 | 适用场景 | 优点 | 缺点 |
|---|---|---|---|
| Optimum CLI(optimum-cli export onnx) | NLP模型的常规导出,尤其是Seq2Seq | 自动处理Encoder-Decoder结构、past key values,几乎开箱即用 | 依赖库版本较新,自定义模型支持有限 |
| torch.onnx.export手写 | 调试、理解模型结构 | 灵活性高,算子映射看得见 | 要自己处理动态轴和past key values,容易漏 |
| onnxruntime.transformers优化 | 针对Bert/GPT类Transformer再优化 | 能自动做fast greedy和beam search优化 | 对Marian等Seq2Seq支持不完整,我试过几次都失败 |
我用的Optimum CLI。原因很直接:MarianMT的导出牵扯到两个子模型(encoder和decoder)、decoder里past key values的输入输出形状,以及动态序列长度。这些细节Optimum已经打磨过了,没必要自己重新趟一遍。
安装依赖我直接给一个能跑通的最小集:
pip install transformers==4.46.2 optimum[exporters]==1.22.2 onnx onnxruntime sentencepiece注意版本不要随便用最新的,我碰到过一次transformers 4.48配合optimum 1.20导出时,pytorch解析模型的返回类型变了,导致导出脚本报"tuple object has no attribute 'last_hidden_state'"。锁定上面这组版本,至少在我这个场景里稳定。
2.2 执行导出命令:一步出两个文件
optimum-cli export onnx --model Helsinki-NLP/opus-mt-en-zh --task translation-en-to-zh translator_onnx/跑完后的目录长这样:
translator_onnx/ ├── config.json ├── decoder_model.onnx ├── encoder_model.onnx ├── generation_config.json ├── sentencepiece.bpe.model ├── special_tokens_map.json ├── tokenizer.json └── tokenizer_config.json看到没有,Optimum自动拆出了encoder和decoder两个文件。这是因为生成式模型在推理时,encoder只需要对输入跑一次,而decoder是循环执行的,两者放在同一个图里会让每次解码都重复跑encoder,性能直接崩。拆开之后,你可以先跑encoder保存结果,再循环跑decoder。
导出时有个动态轴参数值得注意。你如果不加任何参数,导出的模型输入是固定长度。需要在命令里显式声明:
optimum-cli export onnx --model Helsinki-NLP/opus-mt-en-zh --task translation-en-to-zh \ --dynamic-batch-size --dynamic-sequence-length \ translator_onnx/不加这个,你部署时只能处理固定token数的句子,稍微长一点就直接报维度错误。加上之后,输入input_ids的shape会带动态轴,代价是部分runtime优化会被削弱,速度上会差一点点。实际部署时,我建议导出动态模型。
2.3 验证导出结果:图和模型结构要对得上
导出完成后,用onnx.checker检查一下图的合法性,再用下面这段小代码看看输入输出长什么样:
import onnx for model_name in ["encoder_model", "decoder_model"]: model = onnx.load(f"translator_onnx/{model_name}.onnx") print(f"=== {model_name} ===") for inp in model.graph.input: print(f"INPUT: {inp.name}, shape: {[d.dim_value for d in inp.type.tensor_type.shape.dim] if inp.type.tensor_type.shape.dim else 'dynamic'}") for out in model.graph.output: print(f"OUTPUT: {out.name}")正常情况下,encoder_model的输入是input_ids和attention_mask,输出是last_hidden_state之类的encoder对输入序列的语义表示;decoder_model的输入则会多出encoder_outputs、past_key_values等,输出是logits和新的past_key_values。如果格式对不上,就要回到步骤1去检查opset。
3. 验证逻辑:不能只看“翻译出来挺像”
3.1 对照测试:概率分布也要对齐
很多人验证迁移是否成功,就看能不能翻译出一句“意思差不多”的话。这种验证很危险,因为很可能遇到两个模型都输出同一句翻译,但对其他候选词的概率差了一大截,真跑到beam search解码时结果就分岔了。我自己的标准是逐token、逐候选词的概率分布对齐。
写一个对照脚本:
import torch import onnxruntime as ort from transformers import MarianMTModel, MarianTokenizer model_name = "Helsinki-NLP/opus-mt-en-zh" pt_model = MarianMTModel.from_pretrained(model_name) pt_model.eval() tokenizer = MarianTokenizer.from_pretrained(model_name) # ONNX Runtime session so = ort.SessionOptions() so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC enc_session = ort.InferenceSession("translator_onnx/encoder_model.onnx", sess_options=so) dec_session = ort.InferenceSession("translator_onnx/decoder_model.onnx", sess_options=so) text = "This is a simple test for model conversion." enc = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128) with torch.no_grad(): pt_out = pt_model.generate(**enc, num_beams=1, do_sample=False) # 记录PyTorch解码路径中每一步的logits(手动循环更好,这里用generate简化) # ONNX推理 input_ids = enc["input_ids"].numpy() attention_mask = enc["attention_mask"].numpy() encoder_outputs = enc_session.run( ["last_hidden_state"], {"input_ids": input_ids, "attention_mask": attention_mask} )[0] # decoder首次输入 decoder_input_ids = [[tokenizer.eos_token_id]] out_logits = dec_session.run( ["logits"], {"input_ids": decoder_input_ids, "encoder_outputs": encoder_outputs, "past_key_values": []} )[0]实际做这种对比时,要手动实现完整的greedy解码,不让transformers.generate做太多封装,才能在每一步拿到logits。把两个模型的每一步logits做argmax,记录生成的token序列,再计算概率分布的KL散度。如果KL散度在3%以内,说明转换保真度不错;超过5%就要怀疑导出时某个算子精度丢了,优先怀疑Attention里的Softmax或者LayerNorm。
3.2 手动实现greedy解码:这也是部署需要的底层能力
既然要验证,干脆把部署时要用到的greedy解码写出来,一举两得。核心逻辑是:
- encoder只跑一次,得到encoder_outputs。
- decoder每次输入是当前已生成的token序列,加上encoder_outputs和上一步的past_key_values缓存。
- 取logits最后一个token位置,softmax后取argmax,加入生成序列。
- 更新past_key_values,循环直到生成eos或者达到最大长度。
代码片段:
import numpy as np MAX_LEN = 128 past_key_values = [] generated = [tokenizer.eos_token_id] encoder_outputs = enc_session.run(["last_hidden_state"], { "input_ids": input_ids, "attention_mask": attention_mask })[0] for step in range(MAX_LEN): decoder_input = np.array([[generated[-1]]], dtype=np.int64) feeds = { "input_ids": decoder_input, "encoder_outputs": encoder_outputs, "past_key_values": past_key_values, } out = dec_session.run(None, feeds) logits = out[0] # shape: [1, 1, vocab_size] next_token = int(np.argmax(logits[0, -1, :])) generated.append(next_token) past_key_values = out[1:] if next_token == tokenizer.eos_token_id: break print(tokenizer.decode(generated, skip_special_tokens=True))注意这里past_key_values是按层、按方向打包的list of ndarrays,Optimum导出的decoder模型输入输出格式在最上面那个验证脚本里可以打印出来。第一次推理时传空list即可,后面每次把上一次输出的后半部分原样传回去。
我踩过最大的坑是第一次调用时把past_key_values传成了None,导致ONNX Runtime直接崩掉。正确做法是传空列表,让模型知道“没有缓存”。
3.3 回归测试集:准备20句不同类型的中英文对照
不要拿一句“hello world”就敢上线。我整理了一个小测试集,专门覆盖不同难度:
- 陈述句、疑问句、感叹句各5句
- 主语是名词、代词、复合短语各3句,确保代词指代不出错
- 包含数字、日期、专有名词
- 长度从5个词到40个词不等
每一个句子跑完PyTorch和ONNX两条路径,记录生成的token序列是否完全一致,再人工看一遍ONNX输出是否通顺。有两类不一致可以放宽:句子末尾标点差异、数字的“一二三”和“123”之间的差异(这个有时候是模型本身随机性导致,但概率极低)。
4. 性能优化与INT8量化:从能跑到跑得快
4.1 先看baseline:PyTorch vs ONNX原始表现
在没做任何优化前,我分别跑了同一批数据,硬件是i5-10400(6核12线程)、16GB内存、Linux on WSL2的环境,模型opus-mt-en-zh,输入长度大约30个token,输出长度大约20个token。数据如下:
| 推理方式 | 单条生成耗时 | CPU峰值内存 | 模型文件大小 |
|---|---|---|---|
| PyTorch (CPU, greedy) | 约480ms | 约1.1GB | 约300MB(权重fp32) |
| ONNX (CPU, greedy, fp32) | 约230ms | 约650MB | encoder 110MB + decoder 180MB |
| ONNX + INT8动态量化 | 约150ms | 约480MB | encoder 28MB + decoder 46MB |
这个数据不绝对,但趋势非常明显:ONNX直接把延迟压了一半还多,内存也降了;再量化一轮,速度和体积都逼近“能上线”的水准。很多人问为什么ONNX Runtime比PyTorch快这么多,我上面也提过,核心是算子融合和内存规划,这里不再展开。
4.2 动态量化:一行代码做完INT8
ONNX Runtime的量化分动态量化(Dynamic Quantization)和静态量化(Static Quantization)。对NLP推理来说,动态量化默认就够用——它是在推理时实时计算量化的scale和zero point,不用准备校准数据。静态量化需要提前用一批代表性数据跑一遍模型来统计激活值的分布,精度可以更稳,但流程复杂很多。我走的是动态量化:
from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_input="translator_onnx/decoder_model.onnx", model_output="translator_onnx/decoder_model_int8.onnx", weight_type=QuantType.QUInt8, )这里我只对decoder做了量化。encoder模型是一次性前向,不参与自回归循环,量化的收益远没有decoder大。而且encoder做量化后,输出的编码向量误差会被decoder的每一步放大,翻译质量往往差得更明显,所以我建议只量decoder。
量化之后跑一遍3.3节的测试集,BLEU分数(用sacrebleu评测)大概从原始48.2掉到47.5,掉幅在1个点以内,这是可以接受的。如果掉幅超过2个点,换per_channel=True试试(有些权重分布差异大,按通道量化能保住精度)。
4.3 Runtime配置:线程数和内存策略一样重要
量化完了,SessionOptions这层还能再榨出不少性能。我最终的配置长这样:
so = ort.SessionOptions() so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL so.intra_op_num_threads = 4 so.inter_op_num_threads = 1 # 启用内存池复用机制 so.enable_mem_pattern = True so.enable_cpu_mem_arena = True解释一下每个值:
- intra_op_num_threads这个参数控制在单个算子内部的并行度。transformer的MatMul是典型的高并行算子,设4基本能跑满常见CPU的物理核数。注意不要设为12,多线程调度开销反而会让单条延迟变大。
- inter_op_num_threads控制不同算子之间的并行度,也就是pipeline并行,设1就够。如果你的部署并发很大,反过来设4-6更好,可以让多个请求共享等待时间。
- enable_mem_pattern / enable_cpu_mem_arena打开后能复用缓存内存区块,尤其是在多轮长文本生成时,避免每轮解码都重复申请一大块内存。
线程参数需要针对你的实际服务器做交叉验证。我的经验法则是先压测三组配置,对比单条延迟和吞吐,不要照抄别人的数字。
4.4 部署形态:从Python封装到C++/移动端
ONNX迁移的最大价值在于它让模型可以脱离Python环境部署。我在做完Python接口验证后,把同一个ONNX文件接到C++侧直接跑,代码量少得惊人,关键是推理逻辑和Python侧完全一致,只是需要自己管理past_key_values缓存。移动端也有现成的方向,ONNX Runtime的Mobile版本可以把整个模型塞进iOS/Android的包里运行,外人看不出这背后是个大模型。
到了这一步,之前所有在Python侧的验证工作才有意义——因为模型文件已经是一个标准格式,不依赖任何特定框架,你的接口越干净越好迁移。
5. 常见问题与排查技巧实录
我把自己踩过、以及身边朋友踩过的坑整理成一张速查表,可以直接收藏。
5.1 典型问题速查表
| 现象 | 可能原因 | 解决办法 |
|---|---|---|
| HuggingFace模型加载卡住或404 | 网络环境无法稳定访问外网模型仓库 | 提前用huggingface_hub的snapshot_download把模型文件一次性下载到本地目录,之后加载时用local_files_only=True直接读本地文件,不要再走在线协议 |
| optimum-cli导出时报"torch.onnx.export failed: Unsupported operator" | 模型里出现了新算子,不符合onnx标准opset规范 | 升级onnx和onnxscript到最新版;确认opset版本是否过老,可以用--opset 14等参数手动指定 |
| 导出后encoder输出shape为静态,输入稍长就报错 | 没开动态轴 | 重新导出,加上--dynamic-batch-size --dynamic-sequence-length |
| ONNX推理输出和PyTorch生成结果差很远 | 手动解码时past_key_values处理不对,或attention_mask更新漏了 | 打印每一步的logits对比,把attention_mask的生成逻辑补上;先跑greedy模式,不要直接调beam search |
| 同一个模型文件在A机器快、B机器慢 | 线程数、CPU指令集(AVX512、AMX)差异 | 分别在目标机器上跑一遍线程数网格搜索,确认打开CPU指令集优化 |
| INT8量化后输出明显乱码 | attention模块在低精度下数值敏感 | 先换per_channel量化;再不行回退到per-tensor,或者对encoder不量化、只量化decoder |
| decoder第一步传past_key_values为空列表时报错 | ONNX Runtime要求明确的graph输入 | 查graph的输入列表,看看是past_key_values还是past_key_values.1这种实际名字,按名字传入 |
5.2 两个容易被忽略的细节
第一个是tokenizer文件。很多人转换完模型,只盯着onnx文件能跑,但推理时还要用原来的MarianTokenizer做pre/post processing——尤其是它依赖的sentencepiece.bpe.model文件。我导出目录里这个名字很容易被当成普通文件忽略,但一旦缺失,tokenizer直接报FileNotFoundError。部署时记得把整个translator_onnx目录一起拷走,不要只拷两个onnx。
第二个是attention_mask。当处理batch中的短句和长句时,attention_mask不能随便全1,必须保持原始pad位置为0。如果只写了input_ids不考虑mask,短句会被当成“后半段全是正文”,翻译质量直接下降,尤其对batch推理来说影响更明显。
5.3 beam search要不要自己写
很多人迷惑ONNX Runtime导出后,怎么没有beam search功能了。答案是你得自己写。Optimum导出的是模型图,不是生成策略。PyTorch里transformers.generate一行就能beam search,但ONNX Runtime只提供最底层的解码能力。
如果产品上必须要beam search,有两条路:
- 在Python层面实现beam search控制流,每次调用decoder_model获取候选logits,自己维护多个beam路径和past_key_values。参考huggingface generate源码里的beam search部分,代码量大概两三百行,不是不能写。
- 用更轻的采样策略,比如temperature sampling + top_k + top_p,在大部分场景下效果和beam search差不多,但实现简单太多。我自己线上的版本就是采样式的,延迟比beam search低20%左右,用户感知差别很小。
先greedy上线,后面再逐步加复杂采样都是可以的,ONNX图结构根本不用改,只改控制流代码就行,这也是拆encoder/decoder的好处之一。
5.4 模型升级时的回归陷阱
最后提醒一句:ONNX模型不是转完就一劳永逸。我后来更新过HuggingFace上默认的opus-mt版本,结果发现新版模型结构没变,但tokenizer和词表大小变了,导致旧ONNX文件的logits维度和新tokenizer的vocab_size对不上,推理时直接数组越界。所以模型一升级,整套导出、验证、量化流程必须重跑一遍。我给自己的工作流里加了一步CI检查:每次升级模型版本,自动跑一遍20句回归测试,如果BLEU差异超过阈值就告警,避免线上悄悄变差。
6. 写在最后的一些实际经验
迁移ONNX这件事,看起来是导出工具链一步到位,实际上难点全在细节。我最大的体会是“跑通只是开始,量化、线程配置、解码策略、缓存管理,每一层都有优化的空间”。很多人把模型迁到ONNX后跑了个demo就完事了,但ONNX Runtime和PyTorch在真实场景下的表现差异要比demo显现得大得多——多跑几轮压力测试,你会发现在batch、并发、长文本这些指标上,ONNX的领先优势会进一步放大。
另外一个经验是:不要拿到模型就直接怼到外网下载,网络不稳定的环境让你卡在半路很浪费时间。我现在的习惯是先保证本地有完整模型文件,再开始做导出和验证,脱离网络做开发,专注度会好很多。
最后分享一个小技巧:如果你打算在服务端长期部署ONNX模型,可以考虑把encoder和decoder两个模型做成单独的服务,encoder部分可以多路复用给多个decoder线程。因为encoder对同一输入只算一次,之后所有beam和采样试图都共用同一个encoder_outputs,能省下大量重复计算。这个在PyTorch里实现起来比较麻烦,在ONNX Runtime里反而是很自然的工程手法。