1. 为什么要把英译中模型从 HuggingFace 搬到 ONNX1.1 一个真实的需求场景去年年底我接了个离线翻译的小活儿需求很明确在一台没有独立显卡的工控机上跑英译中输入是一段段英文技术文档输出中文要求单句延迟控制在 300ms 以内而且整机不能联网。第一反应当然是去 HuggingFace 上找现成的英译中模型Helsinki-NLP/opus-mt-en-zh这个系列几乎是默认答案质量够用、体积适中、社区验证充分。问题出在部署环节。这台工控机装的是精简版系统Python 环境能跑但 PyTorch 的运行时依赖太重光torch加transformers一套下来磁盘占用就上 G启动加载模型要十几秒内存峰值也压不下来。更麻烦的是推理速度CPU 上跑原始 PyTorch 模型单句要 600ms 往上完全达不到指标。这时候把模型转成 ONNX 就成了最自然的出路。ONNX 的本质是一套开放的模型表示格式它把模型的计算图固化下来脱离训练框架交给专门的推理引擎去跑。转完之后你可以用 ONNX Runtime 加载CPU 上的推理速度通常能提升 2 到 4 倍内存占用也明显下降而且部署时不再需要拖着一整套 PyTorch。所以这篇东西就是把我当时踩过的坑、验证过的步骤完整梳理一遍。适合谁看如果你手上有一个 HuggingFace 上的翻译模型或者任何 encoder-decoder 结构的模型想把它搬到 ONNX 上做轻量化部署那这篇基本可以照着抄。不需要你精通 PyTorch 源码但得能看懂 Python会用 pip 装包知道什么是 tokenizer。1.2 先搞清楚迁移到底在迁什么很多人一上来就torch.onnx.export然后发现导出的模型推理结果不对或者干脆报错。根本原因是没有想清楚一件事HuggingFace 的模型不是一个单纯的计算图它是一整套封装。一个完整的翻译模型至少包含四部分模型权重、模型结构定义forward里的计算逻辑、tokenizer负责文本和 token id 之间的转换、以及生成逻辑beam search、贪心解码这些。ONNX 能固化的只有前两部分里的计算图tokenizer 和生成逻辑是没法直接塞进 ONNX 的。这就引出了迁移的核心思路把模型的计算图导出成 ONNXtokenizer 和生成循环留在 Python 侧用 ONNX Runtime 手动实现。对于 encoder-decoder 结构的翻译模型通常要导出两个 ONNX 文件——一个 encoder一个 decoder带不带 past key values 又是另一个决策点后面细说。理解这一点非常关键因为它决定了你后面所有工作的边界。你不可能导出一个端到端输入英文输出中文的 ONNX 文件然后一劳永逸生成过程必须由你自己控制。想明白这个后面的坑就少踩一半。1.3 ONNX 相比原生 PyTorch 到底赢在哪我把当时实测的数据摆出来环境是 Intel i5-8250U 四核8G 内存无独显模型是opus-mt-en-zh测试集是 200 句平均长度 25 词的英文句子。指标PyTorch 原生ONNX Runtime提升幅度单句平均延迟620ms210ms约 2.9 倍内存峰值1.8GB620MB约 65% 下降模型磁盘占用约 300MB含框架约 90MB约 70% 下降冷启动时间12s2.5s约 4.8 倍延迟的下降主要来自两块一是 ONNX Runtime 对算子做了大量图优化比如算子融合、常量折叠二是它针对 CPU 做了专门的指令集优化能吃到 AVX2 甚至 AVX512。内存下降则是因为不再需要加载整个 PyTorch 运行时。注意上面的数字是特定硬件和模型下的结果换机器、换模型会有差异但量级上的优势是普遍存在的。别把具体数字当承诺把它当参考。2. 动手前的环境准备与模型选型2.1 依赖安装版本对齐是第一道坎这一步看着简单实际上是最容易翻车的地方。transformers、torch、onnx、onnxruntime这四个包的版本之间存在微妙的兼容关系尤其是transformers和torch之间版本差太多会在导出时直接报算子不支持。我当时的组合是这样的实测稳定pip install torch2.1.0 pip install transformers4.35.0 pip install onnx1.15.0 pip install onnxruntime1.16.3 pip install sentencepiece0.1.99几个要点解释一下。sentencepiece必须装因为opus-mt系列的 tokenizer 是基于 sentencepiece 的不装的话加载 tokenizer 会直接失败。onnx和onnxruntime是两个不同的包前者负责导出和校验后者负责推理别搞混。onnxruntime还分 CPU 版和 GPU 版工控机场景装 CPU 版就行包名就是onnxruntimeGPU 版是onnxruntime-gpu。如果你在国内网络环境下装包慢可以配置 pip 的国内源这个属于常规操作配一次省很多事。模型下载同理HuggingFace 的模型仓库在国内访问有时候不稳定可以提前把模型文件下载到本地用本地路径加载这样导出过程完全不依赖网络。提示导出前先用python -c import torch, transformers, onnx, onnxruntime; print(torch.__version__, transformers.__version__)确认版本别等到导出报错了才回头查。2.2 模型选型不是所有翻译模型都好导opus-mt系列是我最推荐的入门选择原因是它的结构标准、社区导出案例多、坑基本都被踩过了。这个系列基于 MarianMT 架构标准的 encoder-decoder注意力机制也是常规实现ONNX 导出支持得很好。如果你用的是 T5 或者 mBART 这类模型导出会复杂一些因为它们的 decoder 结构里有一些动态控制流ONNX 对动态控制流的支持一直是个痛点。不是说不能导而是需要额外处理比如固定 beam size、禁用某些动态特性。选型的时候还要考虑模型大小。opus-mt-en-zh大概 300MB 左右导出成 ONNX 后 fp32 精度约 90MB量化成 int8 后能压到 25MB 左右。如果你对体积敏感量化是必选项但量化会带来轻微的质量损失这个后面单独讲。2.3 目录结构规划别把文件堆一地我习惯在动手前先把目录规划好不然后面文件一多就乱。推荐的结构是这样project/ ├── models/ │ ├── hf_model/ # 原始 HuggingFace 模型 │ └── onnx/ # 导出的 ONNX 文件 │ ├── encoder.onnx │ ├── decoder.onnx │ └── decoder_with_past.onnx ├── scripts/ │ ├── export_onnx.py # 导出脚本 │ └── infer_onnx.py # 推理脚本 └── test_data/ └── samples.txt # 测试句子把原始模型和导出产物分开好处是导出失败可以随时重来不会污染原始文件。测试数据单独放方便做回归对比——每次改完导出参数跑一遍测试集对比输出是否一致这是保证质量的基本功。3. 核心导出流程从 PyTorch 到 ONNX 的完整实操3.1 导出 encoder相对简单但有讲究encoder 的导出是最 straightforward 的部分因为它就是一个标准的 Transformer 编码器输入是 token ids 和 attention mask输出是隐藏状态。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() # 构造 dummy input dummy_text This is a test sentence for export. inputs tokenizer(dummy_text, return_tensorspt) input_ids inputs[input_ids] attention_mask inputs[attention_mask] # 导出 encoder torch.onnx.export( model.get_encoder(), (input_ids, attention_mask), models/onnx/encoder.onnx, input_names[input_ids, attention_mask], output_names[last_hidden_state], dynamic_axes{ input_ids: {0: batch, 1: sequence}, attention_mask: {0: batch, 1: sequence}, last_hidden_state: {0: batch, 1: sequence}, }, opset_version14, do_constant_foldingTrue, )几个关键点必须解释清楚。model.eval()一定要调否则 dropout 层会处于训练模式导出的图里会带上随机性推理结果每次都不一样。dynamic_axes是重中之重它告诉 ONNX 哪些维度是动态的。这里把 batch 和 sequence 两个维度都设成动态意味着导出的模型可以接受任意长度的输入而不是被固定死在 dummy input 的长度上。如果不设你导出的模型就只能处理那一个特定长度的句子完全没法用。opset_version选 14 是个比较稳妥的选择太低的版本不支持某些算子太高的版本可能 onnxruntime 还没跟上。do_constant_foldingTrue让导出时做常量折叠优化能减小模型体积、提升推理速度。3.2 导出 decoderpast key values 是核心难点decoder 的导出是整个流程里最绕的部分绕就绕在past key values这个机制上。翻译模型生成中文的时候是一个词一个词往外蹦的。每生成一个新词decoder 都要重新计算一遍注意力。如果不做优化每步都要把之前所有已生成的词重新算一遍复杂度是 O(n²)。past key values 的作用就是把之前算过的 key 和 value 缓存下来每步只算新词的部分复杂度降到 O(n)。这就导致 decoder 有两种导出方式一种是不带 past 的每步输入完整的已生成序列另一种是带 past 的每步只输入新词加上缓存的 kv。前者简单但慢后者快但导出复杂。生产环境肯定选后者。# 导出带 past key values 的 decoder # 需要构造符合要求的 dummy past import torch num_layers model.config.decoder_layers num_heads model.config.decoder_attention_heads d_model model.config.d_model head_dim d_model // num_heads # 构造 dummy 的 encoder 输出和 past batch_size 1 encoder_seq_len input_ids.shape[1] decoder_seq_len 1 encoder_hidden_states torch.randn(batch_size, encoder_seq_len, d_model) # past key values 的形状是 [batch, num_heads, past_len, head_dim] past_key_values tuple( ( torch.randn(batch_size, num_heads, decoder_seq_len, head_dim), torch.randn(batch_size, num_heads, decoder_seq_len, head_dim), ) for _ in range(num_layers) ) decoder_input_ids torch.tensor([[tokenizer.pad_token_id]])构造 dummy past 的时候层数、头数、维度都必须和模型配置严格对应错一个数字导出就会失败或者推理结果错乱。model.config里能查到这些值别硬编码用配置读出来最保险。导出的时候 input_names 和 output_names 要仔细命名因为后面推理脚本要按名字取输入输出。past 的输入输出是成对的命名上建议用past_key_values.{i}.decoder.key这种带索引的形式方便循环处理。3.3 导出参数逐项拆解导出脚本里那一堆参数每一个都有它的道理我逐个说。dynamic_axes前面说了是让模型支持变长输入的关键。对于 decoder除了 batch 和 sequencepast 的长度维度也要设成动态因为生成过程中 past 会越来越长。opset_version我选 14是因为这个版本对 Transformer 相关算子的支持比较完整尤其是Attention相关的融合算子。如果你导出报算子不支持可以试着降到 12 或 11但可能会损失一些优化。do_constant_folding建议开它会把图里能提前算出来的常量部分算好减小模型体积。export_paramsTrue是默认值表示把权重也写进 ONNX 文件这个必须开不然导出的就是个空壳。use_external_data_format这个参数在模型超过 2GB 的时候需要开因为 protobuf 单文件有 2GB 限制。opus-mt这种小模型用不上但如果你导大模型记得开这个它会把权重单独存成外部文件。3.4 导出后的校验别跳过这一步导出完不校验等于白导。ONNX 官方提供了校验工具能检查图的合法性import onnx model onnx.load(models/onnx/encoder.onnx) onnx.checker.check_model(model) print(ONNX model check passed.)但这只能检查格式合法性不能保证数值正确。真正的校验是拿同一批输入分别跑 PyTorch 和 ONNX对比输出差异。差异在 1e-4 量级以内算正常超过 1e-2 就说明导出有问题。import numpy as np import onnxruntime as ort # PyTorch 输出 with torch.no_grad(): pt_output model.get_encoder()(input_ids, attention_mask).last_hidden_state.numpy() # ONNX 输出 sess ort.InferenceSession(models/onnx/encoder.onnx) onnx_output sess.run( [last_hidden_state], { input_ids: input_ids.numpy(), attention_mask: attention_mask.numpy(), }, )[0] diff np.abs(pt_output - onnx_output).max() print(fMax diff: {diff})这个对比步骤我强烈建议每次都做尤其是你改了导出参数之后。我踩过一次坑改了 opset 版本后输出差异突然变大就是因为某个算子在低版本下实现不同导致数值精度损失。不做对比根本发现不了。4. 推理侧实现用 ONNX Runtime 跑起完整翻译4.1 加载模型与 tokenizer推理侧的第一件事是把 ONNX 模型和 tokenizer 都加载起来。tokenizer 还是用 HuggingFace 的因为它和训练时用的完全一致能保证输入编码不出错。import onnxruntime as ort from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(models/hf_model) # 配置 session 选项 sess_options ort.SessionOptions() sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.intra_op_num_threads 4 encoder_sess ort.InferenceSession( models/onnx/encoder.onnx, sess_optionssess_options, providers[CPUExecutionProvider], ) decoder_sess ort.InferenceSession( models/onnx/decoder_with_past.onnx, sess_optionssess_options, providers[CPUExecutionProvider], )graph_optimization_level设成ORT_ENABLE_ALL能开启所有图优化包括算子融合、内存复用等对性能提升明显。intra_op_num_threads控制单算子内部的并行线程数一般设成物理核心数就行设太大反而会因为线程切换开销导致变慢。providers指定执行后端CPU 场景就是CPUExecutionProvider。如果你有 GPU可以换成CUDAExecutionProvider但要注意 ONNX Runtime 的 GPU 版本需要单独安装而且对 CUDA 版本有要求。4.2 手写贪心解码循环ONNX 不负责生成逻辑所以解码循环得自己写。贪心解码是最简单的策略每步选概率最大的那个 token。import numpy as np def translate(text, max_length128): # 编码输入 inputs tokenizer(text, return_tensorsnp) input_ids inputs[input_ids].astype(np.int64) attention_mask inputs[attention_mask].astype(np.int64) # 跑 encoder encoder_hidden encoder_sess.run( [last_hidden_state], {input_ids: input_ids, attention_mask: attention_mask}, )[0] # 初始化 decoder 输入 decoder_input_ids np.array([[tokenizer.pad_token_id]], dtypenp.int64) past_key_values None generated [] for step in range(max_length): if past_key_values is None: # 第一步没有 past outputs decoder_sess.run( None, { input_ids: decoder_input_ids, encoder_hidden_states: encoder_hidden, }, ) else: # 后续步骤带上 past feed { input_ids: decoder_input_ids, encoder_hidden_states: encoder_hidden, } feed.update(past_key_values) outputs decoder_sess.run(None, feed) logits outputs[0] next_token int(np.argmax(logits[0, -1, :])) generated.append(next_token) if next_token tokenizer.eos_token_id: break # 更新 past 和下一步输入 decoder_input_ids np.array([[next_token]], dtypenp.int64) past_key_values { name: outputs[i 1] for i, name in enumerate(past_output_names) } return tokenizer.decode(generated, skip_special_tokensTrue)这段代码里有几个细节值得说。第一步和后续步骤的输入不一样第一步没有 past后续步骤要带上 past所以循环里有个分支判断。past_output_names是 decoder 输出的 past 对应的名字列表需要和导出时的命名对应上这个得从 ONNX 模型的输出信息里读出来。np.argmax取的是最后一个位置的 logits因为 decoder 每步只预测下一个词。eos_token_id是结束符遇到就停。4.3 从贪心到 beam search 的取舍贪心解码快但质量一般容易陷入局部最优。beam search 保留多个候选质量更好但计算量成倍增加。翻译任务上 beam search 的提升是肉眼可见的尤其是长句。ONNX 侧实现 beam search 会复杂不少因为要维护多个 beam 的 past key values还要处理 beam 之间的合并和排序。我的建议是如果延迟要求不苛刻beam size 设 4 左右质量提升明显如果延迟卡得很死就用贪心或者用 beam size 2 折中。实测下来beam size 4 相比贪心BLEU 大概能提升 1 到 2 个点但延迟增加约 3 倍。这个取舍得根据你的实际场景定。4.4 批处理吞吐量的关键单句推理延迟再低吞吐量上不去也没用。批处理是提升吞吐量的核心手段。ONNX 导出时 batch 维度设成了动态所以天然支持批处理。批处理的关键是 padding。一个 batch 里的句子长度不一要 pad 到同一长度同时用 attention mask 标记哪些是真实 token、哪些是 padding。tokenizer 的batch_encode_plus能自动处理这些。texts [Hello world., This is a longer sentence for testing.] inputs tokenizer(texts, return_tensorsnp, paddingTrue, truncationTrue)批处理下解码循环要稍微改一下因为不同样本可能在不同步数结束。简单做法是跑到所有样本都结束或者达到 max_length对已结束的样本用 eos 填充。复杂做法是动态移除已完成的样本但实现起来麻烦收益有限。5. 量化与优化把模型压到极致5.1 int8 动态量化实操fp32 的 ONNX 模型体积和速度都还有优化空间int8 量化是性价比最高的手段。ONNX Runtime 提供了动态量化工具不需要校准数据集直接就能转。from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_inputmodels/onnx/encoder.onnx, model_outputmodels/onnx/encoder_int8.onnx, weight_typeQuantType.QInt8, )动态量化的原理是把权重从 fp32 压成 int8激活值在推理时动态量化。这样模型体积能压到原来的四分之一左右推理速度也能提升因为 int8 的矩阵乘法比 fp32 快。但量化不是没有代价的。实测下来int8 量化后 BLEU 会掉 0.5 到 1 个点具体取决于模型和测试集。如果你的场景对质量极其敏感可以只量化 encoderdecoder 保持 fp32这样质量损失小一些体积也能降一部分。注意量化后的模型一定要重新跑一遍质量对比别想当然认为应该差不多。我见过量化后输出直接乱码的情况原因是某些层的数值范围超出了 int8 的表示能力这时候需要做 per-channel 量化或者跳过那些层。5.2 图优化与算子融合ONNX Runtime 在加载模型时会自动做图优化但有些优化需要手动开启或者调整。除了前面说的ORT_ENABLE_ALL还可以通过optimized_model_filepath把优化后的模型存下来下次直接加载优化版省去优化时间。sess_options.optimized_model_filepath models/onnx/encoder_optimized.onnx算子融合是图优化的重头戏比如把MatMul Add Gelu融合成一个算子减少内存访问和 kernel 启动开销。这些优化对 Transformer 模型效果尤其明显因为 Transformer 里全是这种可融合的模式。5.3 线程与内存调优CPU 推理的性能和线程配置关系很大。intra_op_num_threads控制算子内并行inter_op_num_threads控制算子间并行。对于 Transformer 这种算子间依赖强的模型inter_op_num_threads设成 1 就行设大了反而因为同步开销变慢。内存方面ONNX Runtime 默认会做内存复用把中间张量的内存池化。如果内存实在紧张可以开启enable_mem_patternFalse牺牲一点速度换内存。但这个开关一般不用动默认配置已经够好。6. 常见问题与排查实录6.1 导出阶段的高频报错导出阶段的问题基本集中在算子支持和形状不匹配两类。我整理了一个速查表报错信息原因解决方式Unsupported operatoropset 版本太低提高 opset_version 到 14 或以上Shape mismatchdummy input 形状和模型期望不符检查 input_ids 和 attention_mask 的维度RuntimeError: expected scalar type输入 dtype 不对确保输入是 int64不是 int32Exporting past_key_values failedpast 结构构造错误核对层数、头数、head_dimModel size exceeds 2GB单文件超限开启 use_external_data_formatUnsupported operator是最常见的尤其是用了比较新的模型结构时。解决办法要么提高 opset要么把那个算子替换成 ONNX 支持的等价实现。有时候需要改模型源码把不支持的算子拆开。6.2 推理结果不对怎么查推理结果不对排查顺序应该是先查 tokenizer再查 encoder最后查 decoder。tokenizer 的问题最常见比如 padding 方向不对、特殊 token 没加。验证方法很简单把 tokenizer 编码再解码看能不能还原原文。还原不了就是 tokenizer 的问题。encoder 的问题用前面说的数值对比法查PyTorch 和 ONNX 输出差异大就说明导出有问题。decoder 的问题最隐蔽因为涉及 past 的传递。常见错误是 past 的维度顺序搞反了或者某一步忘了更新 past。排查方法是把 beam size 设成 1、max_length 设成 3手动打印每一步的输入输出一步步对。6.3 性能不达标的调优思路性能不达标先定位瓶颈在哪。用onnxruntime的 profiling 功能能拿到每个算子的耗时sess_options.enable_profiling True # 跑几次推理后 prof_file sess.end_profiling()打开 profile 文件看哪个算子耗时最多。如果 encoder 耗时占比高说明输入序列太长考虑截断或者分块。如果 decoder 耗时占比高说明生成步数太多考虑调小 max_length 或者优化解码策略。还有一个容易被忽略的点是首次推理的预热。ONNX Runtime 第一次跑会做一些初始化耗时明显偏高。生产环境要在启动后先跑几次空推理预热把初始化开销摊掉。6.4 我踩过的三个坑第一个坑是 dynamic_axes 没设全。当时只设了 batch 维度忘了 sequence 维度结果模型只能处理固定长度输入短句要 pad 到固定长度长句直接报错。这个坑很隐蔽因为导出不报错只有推理时才暴露。第二个坑是 past key values 的 dtype。导出时 past 是 fp32但推理时我传了 fp16 进去结果数值全乱。ONNX 对 dtype 很严格输入输出类型必须完全匹配不能自动转换。第三个坑是量化后没做质量回归。当时图省事量化完直接上线结果用户反馈翻译质量下降明显。后来补做了对比测试发现是 decoder 的某些层量化后精度损失太大改成只量化 encoder 才解决。7. 一些延伸思考与实用建议7.1 模型版本管理别偷懒ONNX 模型一旦导出就和导出时的代码、依赖版本绑定了。建议每次导出都记录模型名称、导出脚本的 git commit、依赖版本、导出参数、校验结果。这些信息在出问题时能救命。我习惯在 ONNX 文件旁边放一个同名的.json元数据文件记录这些信息。看起来麻烦但当你半年后回头要改东西或者要复现某个版本时会感谢当时的自己。7.2 端侧部署的额外考量如果你的目标平台是手机或者嵌入式设备ONNX 可能还不是终点。有些端侧推理框架需要把 ONNX 再转成自己的格式比如转成某些芯片专用的模型格式。这个转换过程又是一轮新的坑主要是算子支持和量化精度的差异。我的建议是先在 PC 上用 ONNX Runtime 把整个流程跑通、验证质量确认没问题了再往端侧转。别一上来就直奔端侧那样出问题你都不知道是导出错了还是转换错了。7.3 什么时候不该用 ONNXONNX 不是万能的。如果你的场景是训练、微调那老老实实用 PyTorchONNX 只适合推理。如果你的模型有大量动态控制流比如某些带条件分支的生成模型ONNX 导出会非常痛苦这时候可以考虑其他推理方案。还有一个情况是模型更新频繁。ONNX 导出是个相对重的流程如果你的模型每周都要更新那维护导出流程的成本可能超过收益。这种场景下直接用 PyTorch 加一些推理优化可能更划算。我个人在实际操作中的体会是ONNX 迁移这件事难点从来不在导出本身而在导出之后的验证和调优。导出脚本网上能搜到一堆但真正决定成败的是你有没有耐心做数值对比、有没有系统地排查问题、有没有在量化后老老实实做质量回归。把这几件事做到位迁移基本就稳了。