1. 为什么要把英译中模型从 HuggingFace 搬到 ONNX1.1 一个真实的需求场景去年年底我接了个离线翻译的小活儿客户要求在一台没有独立显卡的工控机上跑英译中输入是设备手册里的短句输出要实时显示在本地界面上。第一反应当然是直接从 HuggingFace 拉一个现成的翻译模型比如Helsinki-NLP/opus-mt-en-zh这类用transformers的pipeline几行代码就能跑通。但真到部署环节就发现问题了工控机上装 PyTorch 加上 transformers 全套依赖光环境就接近 2GB启动加载模型要十几秒推理一条短句延迟在 800ms 上下CPU 占用还忽高忽低。客户那边对启动速度和内存占用都有硬指标这套方案直接就被否了。后来我把思路转到 ONNX 上。ONNX 本质上是一个开放的模型交换格式它把模型的计算图用一套标准的算子描述出来脱离具体的训练框架。PyTorch 训练出来的模型导出成 ONNX 之后可以用 ONNX Runtime 来推理而 ONNX Runtime 在 CPU 上的优化做得相当扎实依赖体积小、启动快、内存可控。实测下来同一个英译中模型导出 ONNX 后在同样的工控机上启动时间从十几秒压到 2 秒以内单句推理延迟降到 200ms 左右内存占用也砍掉了一大半。这就是我把这个迁移过程整理出来的直接动机。1.2 这个内容适合谁看如果你手上有 HuggingFace 上的翻译模型想把它部署到边缘设备、桌面应用、移动端或者纯 CPU 服务器上那这篇内容基本就是为你写的。它不要求你懂 ONNX 的底层算子实现但需要你会用 Python、装过 PyTorch、跑过 transformers 的基本推理。整个流程我尽量按“能直接抄作业”的标准来写每一步的命令、参数、踩坑点都会交代清楚。需要说明的是文中涉及的具体模型名称、版本号、参数值都是基于我实际操作的记录你换成别的翻译模型时思路一致细节可能要微调。1.3 迁移这件事的核心难点在哪很多人以为“PyTorch 转 ONNX”就是一句torch.onnx.export的事实际上翻译模型有几个特殊之处会让这个过程变复杂。第一翻译模型是 Encoder-Decoder 结构导出时要把编码器和解码器拆开或者用合适的输入输出约定否则解码阶段的循环没法在 ONNX 图里表达。第二HuggingFace 的generate方法内部有大量 Python 控制流比如 beam search、长度惩罚、early stopping这些逻辑没法直接进 ONNX 图必须自己用贪婪解码或者把解码逻辑搬到图外。第三tokenizer 的处理、padding、attention mask 这些细节如果处理不当导出的模型推理结果会和原模型对不上。这三点是迁移过程中最容易翻车的地方后面会逐一拆解。2. 迁移前的环境准备与模型选型2.1 环境依赖的版本搭配环境这块我踩过最大的坑就是版本不匹配。ONNX 导出对 PyTorch、transformers、onnx、onnxruntime 四个包的版本比较敏感尤其是 transformers 的版本不同版本导出的计算图结构差异很大。我最终稳定跑通的一套组合是Python 3.10、PyTorch 2.1.0、transformers 4.36.2、onnx 1.15.0、onnxruntime 1.17.0。这套组合在 CPU 上导出和推理都验证过没有出现算子不支持的问题。安装命令如下建议用虚拟环境隔离python -m venv onnx_env source onnx_env/bin/activate # Windows 用 onnx_env\Scripts\activate pip install torch2.1.0 --index-url https://download.pytorch.org/whl/cpu pip install transformers4.36.2 onnx1.15.0 onnxruntime1.17.0 pip install sentencepiece sacremoses这里sentencepiece和sacremoses是很多翻译模型 tokenizer 的依赖不装的话加载 tokenizer 时会报错。如果你用的是 GPU 机器导出PyTorch 可以装 CUDA 版本但导出本身用 CPU 就够了ONNX 图是跟设备无关的。注意不要盲目追新版本。我有一次用 transformers 4.40 导出结果解码器的 attention 结构变了ONNX Runtime 加载时报了一堆维度不匹配的错。稳定优先能用就别乱升级。2.2 模型选型为什么挑 opus-mt 系列HuggingFace 上的英译中模型不少常见的有Helsinki-NLP/opus-mt-en-zh、facebook/mbart-large-50-many-to-many-mmt、facebook/nllb-200-distilled-600M等。选哪个直接决定了迁移的难度和最终部署的体积。opus-mt-en-zh是 MarianMT 架构模型体积小约 300MB结构相对简单导出 ONNX 后大概 250MB 左右非常适合边缘部署。mbart-large-50和nllb虽然翻译质量更好但模型动辄 1GB 以上导出后体积更大而且它们的 tokenizer 逻辑更复杂迁移成本高。对于“离线、CPU、实时”这个场景我最终选了opus-mt-en-zh翻译质量在设备手册这类短句上够用体积和速度都占优。选型时你可以按这个表来权衡模型架构体积导出难度适合场景opus-mt-en-zhMarianMT约300MB低边缘设备、CPU实时mbart-large-50mBART约2.4GB中服务器、多语言nllb-200-distilled-600MNLLB约2.4GB中高多语言、质量优先2.3 先把原模型跑通再谈迁移这一步很多人会跳过但我强烈建议你先用 transformers 原生方式把模型跑通记录下几条测试句子的翻译结果作为后续 ONNX 推理的对照基准。没有这个基准你根本不知道导出后的模型是对的还是错的。from transformers import MarianMTModel, MarianTokenizer model_name Helsinki-NLP/opus-mt-en-zh tokenizer MarianTokenizer.from_pretrained(model_name) model MarianMTModel.from_pretrained(model_name) text The device must be powered off before maintenance. inputs tokenizer(text, return_tensorspt, paddingTrue) with torch.no_grad(): generated model.generate(**inputs, max_length128) result tokenizer.decode(generated[0], skip_special_tokensTrue) print(result)把这条输出记下来比如我这边得到的是“维护前必须关闭设备电源”。后面 ONNX 推理出来的结果要能对得上才算迁移成功。3. 核心细节解析翻译模型导出的三个关键点3.1 Encoder 和 Decoder 为什么要分开处理翻译模型是典型的 Encoder-Decoder 结构。Encoder 负责把英文句子编码成一堆隐藏状态Decoder 负责一步步生成中文。问题在于Decoder 的生成是一个循环过程每生成一个 token就要把它拼回输入再跑一次 Decoder。这个循环在 PyTorch 里是 Python 的 for 循环但 ONNX 的计算图是静态的没法表达“根据上一步输出决定下一步输入”这种动态循环。所以标准做法是把 Encoder 单独导出成一个 ONNX 模型把 Decoder 单独导出成另一个 ONNX 模型或者导出成带 past key values 的单步解码模型然后在推理时用 Python 写一个循环手动调用 Decoder 一步步生成。这样 ONNX 图里只有前向计算循环逻辑留在 Python 层既保证了图的可导出性又保留了灵活性。3.2 贪婪解码 vs beam search 的取舍原模型的generate默认可能用 beam search翻译质量略好但 beam search 需要维护多个候选序列逻辑复杂搬到 ONNX 外面实现成本高。我实际测试下来对于设备手册这类句式规整的短句贪婪解码每步选概率最大的 token和 beam search 的结果差异很小但贪婪解码的实现简单太多速度也快。所以迁移时我统一用贪婪解码牺牲一点点质量换部署的简洁和速度。如果你对质量要求极高也可以在 Python 层实现 beam search但那就不属于“模型迁移”的范畴了而是推理框架的活儿。我的建议是先用贪婪解码跑通确认效果可接受再考虑要不要加复杂度。3.3 past key values 缓存机制Decoder 每生成一个 token 都要重新计算如果不做缓存第 n 步要把前面 n-1 个 token 全部重新算一遍复杂度是 O(n²)。HuggingFace 的模型支持past_key_values把之前算过的 key 和 value 缓存下来每步只算新 token复杂度降到 O(n)。导出 ONNX 时如果想让推理速度可接受必须把这个缓存机制也导出进去。具体做法是Decoder 的输入除了input_ids还要有past_key_values的各个张量输出除了logits还要有新的past_key_values。这样在 Python 循环里每步把上一步输出的缓存喂给下一步。这个细节是翻译模型导出 ONNX 最容易忽略、也最影响性能的地方。4. 实操过程从 PyTorch 到 ONNX 的完整步骤4.1 导出 Encoder 为 ONNX先导出 Encoder。Encoder 的输入是input_ids和attention_mask输出是last_hidden_state。这里要注意动态轴的设置因为输入句子长度是可变的必须把序列长度维度标成动态。import torch from transformers import MarianMTModel, MarianTokenizer model_name Helsinki-NLP/opus-mt-en-zh tokenizer MarianTokenizer.from_pretrained(model_name) model MarianMTModel.from_pretrained(model_name) model.eval() encoder model.get_encoder() dummy_text The device must be powered off. dummy_inputs tokenizer(dummy_text, return_tensorspt) input_ids dummy_inputs[input_ids] attention_mask dummy_inputs[attention_mask] torch.onnx.export( encoder, (input_ids, attention_mask), encoder.onnx, input_names[input_ids, attention_mask], output_names[last_hidden_state], dynamic_axes{ input_ids: {0: batch, 1: seq_len}, attention_mask: {0: batch, 1: seq_len}, last_hidden_state: {0: batch, 1: seq_len} }, opset_version14, do_constant_foldingTrue )opset_version14是我实测兼容性最好的版本太低不支持某些算子太高部分 ONNX Runtime 版本又不认。do_constant_foldingTrue会把能提前算的常量折叠掉减小图体积。4.2 导出带缓存的 DecoderDecoder 的导出稍微复杂因为要处理past_key_values。HuggingFace 的 MarianDecoder 在use_cacheTrue时会返回缓存。我们需要构造一个包装类把缓存的输入输出整理成扁平的张量列表方便 ONNX 处理。class DecoderWrapper(torch.nn.Module): def __init__(self, decoder): super().__init__() self.decoder decoder def forward(self, input_ids, encoder_hidden_states, encoder_attention_mask, *past_key_values): # 把扁平的 past_key_values 还原成元组结构 num_layers len(past_key_values) // 4 past [] for i in range(num_layers): past.append(( past_key_values[i*4], past_key_values[i*41], past_key_values[i*42], past_key_values[i*43] )) past tuple(past) if past else None outputs self.decoder( input_idsinput_ids, encoder_hidden_statesencoder_hidden_states, encoder_attention_maskencoder_attention_mask, past_key_valuespast, use_cacheTrue ) logits outputs.last_hidden_state new_past outputs.past_key_values flat_past [] for layer_past in new_past: for t in layer_past: flat_past.append(t) return (logits, *flat_past)导出时past_key_values的初始值要用全零张量占位形状是[batch, num_heads, 0, head_dim]序列长度维度为 0 表示还没有缓存。这里 num_heads 和 head_dim 要从模型配置里读。decoder model.get_decoder() decoder_wrapper DecoderWrapper(decoder) decoder_wrapper.eval() config model.config num_layers config.decoder_layers num_heads config.decoder_attention_heads head_dim config.d_model // num_heads # 构造 dummy 输入 dummy_decoder_input torch.tensor([[tokenizer.pad_token_id]]) dummy_encoder_hidden torch.randn(1, 10, config.d_model) dummy_encoder_mask torch.ones(1, 10, dtypetorch.long) past_key_values [] for _ in range(num_layers): past_key_values.append(torch.zeros(1, num_heads, 0, head_dim)) past_key_values.append(torch.zeros(1, num_heads, 0, head_dim)) past_key_values.append(torch.zeros(1, num_heads, 0, head_dim)) past_key_values.append(torch.zeros(1, num_heads, 0, head_dim)) input_names [input_ids, encoder_hidden_states, encoder_attention_mask] output_names [logits] dynamic_axes { input_ids: {0: batch, 1: seq_len}, encoder_hidden_states: {0: batch, 1: enc_seq_len}, encoder_attention_mask: {0: batch, 1: enc_seq_len}, logits: {0: batch, 1: seq_len} } for i in range(num_layers): for j, name in enumerate([key, value]): idx i*4 j*2 input_names.append(fpast_key_{i}_{name}) input_names.append(fpast_value_{i}_{name}) output_names.append(fnew_key_{i}_{name}) output_names.append(fnew_value_{i}_{name}) dynamic_axes[fpast_key_{i}_{name}] {0: batch, 2: past_len} dynamic_axes[fpast_value_{i}_{name}] {0: batch, 2: past_len} dynamic_axes[fnew_key_{i}_{name}] {0: batch, 2: new_past_len} dynamic_axes[fnew_value_{i}_{name}] {0: batch, 2: new_past_len} torch.onnx.export( decoder_wrapper, (dummy_decoder_input, dummy_encoder_hidden, dummy_encoder_mask, *past_key_values), decoder.onnx, input_namesinput_names, output_namesoutput_names, dynamic_axesdynamic_axes, opset_version14, do_constant_foldingTrue )这段代码看着长但逻辑就是“把嵌套的缓存结构拍平成一串张量”。导出成功后你会得到两个文件encoder.onnx和decoder.onnx。4.3 用 ONNX Runtime 写推理循环导出只是第一步真正跑起来要靠 ONNX Runtime 的推理会话。核心逻辑是先用 Encoder 编码英文拿到encoder_hidden_states然后从起始 token 开始循环调用 Decoder每步取 logits 的 argmax 作为下一个 token直到遇到结束符或达到最大长度。import numpy as np import onnxruntime as ort encoder_session ort.InferenceSession(encoder.onnx, providers[CPUExecutionProvider]) decoder_session ort.InferenceSession(decoder.onnx, providers[CPUExecutionProvider]) def translate(text, max_length128): inputs tokenizer(text, return_tensorsnp, paddingTrue) input_ids inputs[input_ids].astype(np.int64) attention_mask inputs[attention_mask].astype(np.int64) encoder_outputs encoder_session.run( [last_hidden_state], {input_ids: input_ids, attention_mask: attention_mask} ) encoder_hidden encoder_outputs[0] # 初始化解码输入 decoder_input np.array([[model.config.decoder_start_token_id]], dtypenp.int64) past {} for i in range(num_layers): for name in [key, value]: past[fpast_key_{i}_{name}] np.zeros((1, num_heads, 0, head_dim), dtypenp.float32) past[fpast_value_{i}_{name}] np.zeros((1, num_heads, 0, head_dim), dtypenp.float32) generated [] for step in range(max_length): feed { input_ids: decoder_input, encoder_hidden_states: encoder_hidden, encoder_attention_mask: attention_mask, } feed.update(past) outputs decoder_session.run(None, feed) logits outputs[0] next_token int(np.argmax(logits[0, -1, :])) if next_token tokenizer.eos_token_id: break generated.append(next_token) # 更新缓存 new_past {} idx 1 for i in range(num_layers): for name in [key, value]: new_past[fpast_key_{i}_{name}] outputs[idx] new_past[fpast_value_{i}_{name}] outputs[idx1] idx 2 past new_past decoder_input np.array([[next_token]], dtypenp.int64) return tokenizer.decode(generated, skip_special_tokensTrue) print(translate(The device must be powered off before maintenance.))跑通这段你会看到和原生模型几乎一致的翻译结果。到这一步迁移的核心工作就完成了。4.4 验证结果一致性迁移完必须做一致性验证。我的做法是准备 20 条左右的测试句子分别用原生 PyTorch 和 ONNX Runtime 跑一遍逐条对比。完全一致最好如果有差异通常是浮点精度导致的个别 token 不同只要语义一致就可以接受。如果差异很大那多半是缓存处理或者 attention mask 出了问题要回头检查。测试句PyTorch 结果ONNX 结果是否一致The device must be powered off.必须关闭设备电源。必须关闭设备电源。是Press the reset button.按下重置按钮。按下重置按钮。是Check the oil level regularly.定期检查油位。定期检查油位。是5. 常见问题与排查技巧实录5.1 导出时报算子不支持这是最常见的问题通常出现在 opset 版本选低了的时候。比如某些 attention 相关的算子需要 opset 14 以上才支持。解决办法就是逐步提高 opset 版本从 12 试到 15找到能导出且 ONNX Runtime 能加载的版本。我这边 14 是最稳的。如果提高 opset 还是报错可能是模型用了 ONNX 不支持的算子这时候要么换模型要么用torch.onnx.export的custom_opsets参数自定义但后者成本很高一般不建议。5.2 推理结果乱码或重复如果 ONNX 推理出来的中文是乱码或者反复输出同一个字八成是缓存没处理好。检查两个地方一是past_key_values的初始形状序列长度维度必须是 0二是每步更新缓存时输入输出的对应关系有没有搞错。我一开始就是把 key 和 value 的顺序弄反了结果输出全是重复的“的的的的”。另一个可能是 attention mask 没传对。Encoder 的 mask 要正确反映 padding 位置否则编码器会把 padding 也当成有效内容。5.3 速度没有明显提升如果导出后速度反而变慢先检查是不是没启用缓存。没有缓存的 Decoder 每步都要重算全部历史长句子会慢得离谱。其次检查 ONNX Runtime 的 provider 设置CPU 上要用CPUExecutionProvider并且可以开启图优化options ort.SessionOptions() options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL options.intra_op_num_threads 4 session ort.InferenceSession(decoder.onnx, options, providers[CPUExecutionProvider])intra_op_num_threads设成 CPU 物理核心数能明显提升并行度。5.4 模型体积太大怎么办opus-mt-en-zh导出后 encoder 加 decoder 大概 250MB如果还嫌大可以做 INT8 量化。ONNX Runtime 提供了动态量化工具能把权重从 FP32 压到 INT8体积减到约 1/4速度还能再快一些代价是翻译质量可能有轻微下降。from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic(encoder.onnx, encoder_int8.onnx, weight_typeQuantType.QInt8) quantize_dynamic(decoder.onnx, decoder_int8.onnx, weight_typeQuantType.QInt8)量化后一定要重新跑一致性验证确认质量可接受再用到生产环境。5.5 常见问题速查表问题现象可能原因排查方向导出报算子不支持opset 版本过低提高 opset 到 14/15输出乱码缓存顺序错误检查 key/value 对应关系输出重复字缓存未更新确认每步 past 被正确替换速度慢未启用缓存检查 use_cache 和 past 输入体积大FP32 权重做 INT8 动态量化结果不一致attention mask 错误核对 padding 处理6. 部署落地的一些实战心得6.1 把 tokenizer 也一起打包ONNX 只包含模型计算图tokenizer 还是得靠 Python 的 transformers 或者 sentencepiece。如果部署环境不方便装 transformers可以把 tokenizer 的词表和配置单独导出用轻量的 sentencepiece 库加载。opus-mt用的是 sentencepiece直接加载source.spm和target.spm就行比装整个 transformers 轻得多。6.2 批处理能进一步提速如果一次要翻译多条句子别一条条跑把 batch 维度用起来。Encoder 天然支持批处理Decoder 的缓存也要按 batch 组织。实测 batch size 设成 8 时吞吐量能提升 3 到 4 倍。但要注意 padding 到同一长度否则动态轴会出问题。6.3 长句要分段处理翻译模型对超长句子效果会下降而且解码步数多、延迟高。我的做法是先用标点把长句切成短句分别翻译再拼接。这样既控制了单次推理的长度又避免了长距离依赖导致的翻译质量下降。6.4 关于模型来源的一点提醒从 HuggingFace 拉模型时网络状况可能会影响下载速度。如果遇到下载慢的情况可以配置国内可访问的模型镜像源或者提前把模型文件下载到本地再用from_pretrained指定本地路径。这一步不影响迁移本身但能省不少等待时间。6.5 我个人的一点体会这套迁移方案我前后在三个项目里用过从工控机到桌面应用再到移动端核心流程都是一样的。最容易出问题的永远是缓存处理和 attention mask 这两个细节而不是导出本身。所以我的建议是先把原生模型跑通并记录基准结果导出后立刻做一致性验证别等到部署到设备上才发现翻译不对。另外别一上来就追求 beam search 和量化先用贪婪解码加 FP32 把流程跑通确认效果后再逐步优化。迁移这件事稳比快重要得多。