1. 为什么要把英译中模型从原生框架搬到 ONNX做英译中这类序列到序列的任务很多人第一反应是直接调用现成的推理接口跑通就行。但真到了要落地的时候问题就来了推理速度上不去、显存占用高、部署环境依赖一大堆、换个平台就得重新配环境。我自己在几个实际项目里都遇到过这种情况尤其是需要把模型塞进边缘设备或者做批量离线翻译的时候原生框架的推理效率实在让人头疼。ONNX 这个东西全称是 Open Neural Network Exchange说白了就是一种中间表示格式。它最大的好处是跨框架、跨平台而且配合 ONNX Runtime 做推理性能通常比原生 PyTorch 直接跑要快不少。我实测下来在 CPU 上做英译中推理ONNX Runtime 的延迟大概能降到原来的六成左右内存占用也能压下来一截。这不是玄学是因为 ONNX Runtime 在算子融合、内存复用、线程调度这些底层做了大量优化。那为什么选 HuggingFace 上的英译中模型原因很简单生态成熟、模型选择多、tokenizer 配套齐全。你不需要自己从头训练一个翻译模型直接拿现成的预训练权重做迁移就行。常见的选择包括基于 Transformer 架构的各类翻译模型比如类似 Marian、T5、BART 这些结构。它们在中英翻译任务上表现稳定社区里也有大量微调好的权重可以直接用。这个迁移过程适合谁来参考我觉得有三类人比较需要第一类是做端侧部署的工程师需要在没有 GPU 的环境里跑翻译第二类是做批量离线处理的比如一次性翻译几万条文本这时候推理效率直接决定成本第三类是想学习模型格式转换的开发者ONNX 导出是个很好的练手项目能帮你理解计算图和算子层面的东西。整个迁移的核心思路其实不复杂加载预训练模型和 tokenizer构造一个带 past key values 的推理包装用 torch.onnx.export 导出计算图然后用 ONNX Runtime 加载并做推理验证。但魔鬼在细节里seq2seq 模型的导出比普通分类模型麻烦得多因为涉及编码器-解码器结构、自回归生成、KV 缓存这些机制。下面我会把每个环节拆开讲包括我踩过的坑和最终跑通的方案。2. 迁移前的环境准备与模型选型2.1 依赖安装与版本匹配环境这块我建议用 Python 3.9 到 3.11 之间的版本太新或太旧都可能遇到 wheel 包不兼容的问题。核心依赖就几个transformers、torch、onnx、onnxruntime。如果你要用 GPU 做导出还得装 onnxruntime-gpu但导出过程本身其实 CPU 就够了。pip install transformers4.36.0 pip install torch2.1.0 pip install onnx1.15.0 pip install onnxruntime1.16.0版本这块我要特别提醒一句transformers 和 torch 的版本要匹配好。我有一次用 transformers 4.40 配 torch 2.0导出的时候直接报算子不支持。后来降到 transformers 4.36 配 torch 2.1问题就没了。原因是新版本 transformers 里有些算子实现变了ONNX 导出器还没跟上。所以别盲目追新稳定组合比最新版本重要。另外onnx 和 onnxruntime 的版本也要对齐。onnx 是定义格式的onnxruntime 是执行引擎两者版本差太多会出现模型能导出但加载失败的情况。我一般用 onnx 1.15 配 onnxruntime 1.16这个组合在我多个项目里都跑得很稳。2.2 英译中模型的选择逻辑HuggingFace 上英译中模型不少怎么选我的经验是看三个维度模型大小、翻译质量、导出友好度。模型大小直接决定推理速度和内存占用。base 级别的模型大概 2 到 3 亿参数导出后的 ONNX 文件在 300MB 到 500MB 之间CPU 上跑单句翻译大概几百毫秒。large 级别的模型参数翻倍质量确实好一些但推理速度会慢很多端侧基本不用考虑。翻译质量这块我建议先在小批量测试集上跑一遍看看 BLEU 或者直接人工看几十条结果。有些模型在通用领域表现好但到了专业领域就拉胯。如果你做的是特定领域的翻译最好找在该领域微调过的权重。导出友好度是最容易被忽略的。有些模型用了自定义算子或者复杂的注意力实现导出 ONNX 的时候会报一堆警告甚至直接失败。我一般优先选那些社区里已经有人成功导出过的模型结构比如标准的 encoder-decoder Transformer这类结构导出路径最成熟。提示选模型的时候先看它的 config.json 里 model_type 是什么。如果是 marian、t5、bart 这些常见类型导出成功率很高。如果是自定义架构建议先做个小规模导出测试再决定。2.3 导出前的模型检查清单在动手导出之前我习惯先做几项检查能省掉后面很多调试时间。第一确认模型能正常做翻译。加载模型和 tokenizer随便输入一句英文看看输出是不是合理的中文。这一步是基线如果原生推理都有问题导出肯定也白搭。第二检查模型的输入输出签名。用model.forward的参数列表看看需要哪些输入。seq2seq 模型通常需要 input_ids、attention_mask解码阶段还需要 decoder_input_ids 和 past_key_values。第三确认 tokenizer 的类型和特殊 token。英译中模型一般用 SentencePiece 或者 BPE特殊 token 包括 pad、eos、decoder_start 这些。导出后的推理需要手动处理这些 token所以必须提前搞清楚。第四测一下原生推理的耗时和内存。这是后面对比 ONNX 性能的基准。我一般用 time.perf_counter 测 100 次取平均内存用 psutil 或者 torch.cuda.max_memory_allocated 来看。3. 核心转换流程的详细拆解3.1 Seq2Seq 模型导出的特殊难点普通分类模型导出 ONNX 很简单一个 forward 就完事。但 seq2seq 模型不行因为它的生成过程是自回归的先编码源句子然后一步步解码每步的输出会作为下一步的输入。这个循环在 PyTorch 里是 Python 代码控制的ONNX 的计算图是静态的没法直接表达循环。解决方案有两种思路。第一种是把整个生成循环也导出成图用 ONNX 的 Loop 算子实现。这种做法理论上最完整但实际非常难搞因为循环里涉及动态长度、条件判断、KV 缓存更新导出器经常处理不好。第二种思路是把模型拆成编码器和解码器两部分分别导出。编码器跑一次得到隐藏状态解码器每次跑一步KV 缓存由外部代码管理。这种做法灵活、可控也是我实际采用并推荐的方式。缺点是需要自己写生成循环但代码量不大而且调试起来方便得多。我选第二种还有一个原因ONNX Runtime 对动态 shape 的支持虽然不错但在自回归场景下每次解码的输入长度都在变如果整个循环都在图里动态 shape 的处理会很复杂。拆开之后编码器处理固定长度的源句子解码器每次处理长度为 1 的新 tokenshape 管理简单很多。3.2 编码器导出实操编码器的作用是把源语言句子编码成一组隐藏状态。输入是 input_ids 和 attention_mask输出是 last_hidden_state。import torch from transformers import AutoModelForSeq2SeqLM, AutoTokenizer model_name your-en-zh-model tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForSeq2SeqLM.from_pretrained(model_name) model.eval() encoder model.get_encoder() dummy_input_ids torch.randint(0, 1000, (1, 16), dtypetorch.long) dummy_attention_mask torch.ones((1, 16), dtypetorch.long) torch.onnx.export( encoder, (dummy_input_ids, dummy_attention_mask), encoder.onnx, input_names[input_ids, attention_mask], output_names[encoder_hidden_states], dynamic_axes{ input_ids: {0: batch, 1: src_len}, attention_mask: {0: batch, 1: src_len}, encoder_hidden_states: {0: batch, 1: src_len} }, opset_version14, do_constant_foldingTrue )这里有几个关键点。opset_version 我选 14因为 seq2seq 模型里的注意力算子在这个版本支持比较好。do_constant_folding 打开能做常量折叠优化减小模型体积。dynamic_axes 必须设置否则导出的模型只能处理固定长度输入那就没用了。导出过程中可能会看到一些警告比如某些算子被拆解或者不支持。大部分警告可以忽略只要最后模型能加载并输出正确结果就行。但如果看到 Exporting operator X is not supported 这种就要注意了可能需要换 opset 版本或者改模型实现。3.3 解码器导出与 KV 缓存处理解码器是重头戏。它每次接收当前 token、编码器输出、以及之前的 KV 缓存输出下一个 token 的 logits 和更新后的 KV 缓存。decoder model.get_decoder() lm_head model.lm_head class DecoderWrapper(torch.nn.Module): def __init__(self, decoder, lm_head): super().__init__() self.decoder decoder self.lm_head lm_head def forward(self, decoder_input_ids, encoder_hidden_states, encoder_attention_mask, *past_key_values): outputs self.decoder( input_idsdecoder_input_ids, encoder_hidden_statesencoder_hidden_states, encoder_attention_maskencoder_attention_mask, past_key_valuespast_key_values if past_key_values else None, use_cacheTrue ) logits self.lm_head(outputs.last_hidden_state) return (logits,) outputs.past_key_valuesKV 缓存的处理是难点。每个解码层都有 key 和 value 两个缓存张量形状是 (batch, num_heads, past_len, head_dim)。导出的时候要把这些缓存作为输入和输出都声明出来。层数多的话输入输出参数会很多但这是必须的。我第一次导出的时候没加 KV 缓存结果每次解码都要重新计算所有历史 token 的注意力速度慢得离谱。加上缓存之后解码速度提升了大概三到四倍。所以这个优化绝对不能省。注意KV 缓存的形状里 past_len 是动态的导出时 dynamic_axes 要把它标成动态维度。另外第一次解码时 past_len 为 0有些版本的导出器处理不了长度为 0 的输入可以用一个长度为 1 的 dummy 缓存代替然后在外部逻辑里忽略它。3.4 导出后的模型验证方法导出完不能直接信必须验证输出和原生模型一致。我的做法是用同一组输入分别跑原生 PyTorch 和 ONNX Runtime对比输出的数值差异。import numpy as np import onnxruntime as ort sess ort.InferenceSession(encoder.onnx) onnx_out sess.run(None, { input_ids: dummy_input_ids.numpy(), attention_mask: dummy_attention_mask.numpy() }) with torch.no_grad(): torch_out encoder(dummy_input_ids, dummy_attention_mask)[0].numpy() diff np.abs(onnx_out[0] - torch_out).max() print(fMax diff: {diff})一般来说max diff 在 1e-4 到 1e-3 之间是正常的因为浮点运算顺序不同会带来微小误差。如果 diff 超过 1e-2那就有问题了可能是某个算子导出不正确需要逐层排查。排查的方法是把编码器拆成更小的模块逐个导出对比。比如先单独导出 embedding 层再导出第一层注意力这样能定位到具体是哪一层出了问题。这个过程比较繁琐但比盲目改代码有效得多。4. ONNX Runtime 推理的完整实现4.1 推理会话的配置与优化ONNX Runtime 的推理会话有很多可配置项配好了性能提升明显。options ort.SessionOptions() options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL options.intra_op_num_threads 4 options.inter_op_num_threads 2 options.execution_mode ort.ExecutionMode.ORT_SEQUENTIAL sess ort.InferenceSession(encoder.onnx, options, providers[CPUExecutionProvider])graph_optimization_level 设成 ORT_ENABLE_ALL 会启用所有图优化包括算子融合、常量折叠、冗余消除。intra_op_num_threads 控制单个算子内部的并行线程数一般设成物理核心数。inter_op_num_threads 控制算子之间的并行度对于 seq2seq 这种串行依赖强的模型设小一点反而好。如果你有 GPU可以把 providers 换成 CUDAExecutionProvider。但要注意GPU 上的 ONNX Runtime 对动态 shape 的支持不如 CPU 稳定有时候需要把输入 padding 到固定长度才能跑。我一般先在 CPU 上调通再尝试 GPU。4.2 自回归生成循环的编写生成循环的逻辑是编码源句子然后从 decoder_start_token 开始每次解码一个 token直到遇到 eos 或者达到最大长度。def translate(text, encoder_sess, decoder_sess, tokenizer, max_len128): inputs tokenizer(text, return_tensorsnp) encoder_hidden encoder_sess.run(None, { input_ids: inputs[input_ids].astype(np.int64), attention_mask: inputs[attention_mask].astype(np.int64) })[0] decoder_input np.array([[tokenizer.pad_token_id]], dtypenp.int64) past_kv None generated [] for _ in range(max_len): feed { decoder_input_ids: decoder_input, encoder_hidden_states: encoder_hidden, encoder_attention_mask: inputs[attention_mask].astype(np.int64) } if past_kv is not None: for i, kv in enumerate(past_kv): feed[fpast_key_{i}] kv[0] feed[fpast_value_{i}] kv[1] outputs decoder_sess.run(None, feed) logits outputs[0] past_kv [(outputs[1 2*i], outputs[2 2*i]) for i in range((len(outputs)-1)//2)] next_token int(np.argmax(logits[0, -1, :])) if next_token tokenizer.eos_token_id: break generated.append(next_token) decoder_input np.array([[next_token]], dtypenp.int64) return tokenizer.decode(generated, skip_special_tokensTrue)这个循环里有个细节第一次解码时 past_kv 是 Nonefeed 里不带缓存。从第二次开始把上一次的输出缓存塞进去。缓存的索引要对应好每个解码层有两个张量key 和 value顺序不能乱。4.3 性能对比与调优记录我在一台 8 核 CPU 的机器上做了对比测试源句子平均长度 20 个 token目标句子平均 25 个 token。方案平均延迟内存占用模型体积PyTorch 原生820ms1.2GB950MBONNX Runtime CPU480ms680MB920MBONNX Runtime CPU 量化310ms420MB240MB量化那行是用 ONNX Runtime 的动态量化做的把权重从 FP32 转成 INT8。精度损失很小我测了几百条翻译BLEU 只掉了 0.3 左右但速度和内存改善非常明显。如果你的场景对精度不是极端敏感强烈建议做量化。量化的命令很简单from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic(encoder.onnx, encoder_quant.onnx, weight_typeQuantType.QUInt8) quantize_dynamic(decoder.onnx, decoder_quant.onnx, weight_typeQuantType.QUInt8)提示量化后的模型在第一次推理时会稍慢因为要做一些初始化。从第二次开始就稳定快了。所以做性能测试的时候要跑多次取平均别只看第一次。5. 常见问题排查与避坑经验5.1 导出阶段的典型报错导出阶段最容易遇到的问题是算子不支持。比如某些模型用了自定义的注意力实现里面有 ONNX 不认识的算子。解决办法通常是改 opset 版本或者把模型里的自定义部分替换成标准实现。另一个常见问题是动态 shape 处理失败。报错信息一般是 Dynamic shape not supported for operator X。这时候要检查 dynamic_axes 的设置确保所有变化的维度都标出来了。特别是 KV 缓存里的 past_len 维度很容易漏掉。还有一种情况是导出成功但输出全为 NaN。这通常是因为模型里有除零或者 log(0) 的操作在导出时被常量折叠算出了 NaN。解决办法是关掉 do_constant_folding或者检查模型里有没有数值不稳定的地方。5.2 推理阶段的精度与速度问题推理阶段最常见的问题是输出和原生不一致。如果差异很大先检查输入数据的类型和形状。ONNX Runtime 对输入类型很严格input_ids 必须是 int64attention_mask 也必须是 int64类型不对会直接报错或者输出垃圾。速度不达预期的话先看线程配置。默认情况下 ONNX Runtime 可能只用单线程性能差很多。把 intra_op_num_threads 设成 CPU 核心数通常能提升两三倍。另外检查有没有启用图优化ORT_ENABLE_ALL 和 ORT_DISABLE_ALL 的性能差距可能有一倍。如果用了 GPU 但速度没提升可能是数据传输成了瓶颈。每次推理都要把输入从 CPU 拷到 GPU输出再拷回来这个开销在小 batch 场景下很致命。解决办法是把整个生成循环都放在 GPU 上或者用 IO binding 减少拷贝。5.3 常见问题速查表问题现象可能原因解决方向导出报算子不支持opset 版本过低或模型用了自定义算子提高 opset 到 14 以上替换自定义实现导出成功但加载失败onnx 和 onnxruntime 版本不匹配对齐两者版本推荐 onnx 1.15 ort 1.16输出全为 NaN常量折叠算出异常值关闭 do_constant_folding输出和原生差异大输入类型或形状不对检查 dtype 是否为 int64shape 是否匹配推理速度慢线程数没配好或没开图优化设置 intra_op_num_threads启用 ORT_ENABLE_ALL第一次推理特别慢量化模型初始化开销预热几次后再测性能GPU 上速度没提升数据传输开销大用 IO binding 或把循环放在 GPU 上5.4 我踩过的几个坑第一个坑是忘了设置 model.eval()。训练模式下的 dropout 和 batch norm 行为不一样导出后的模型输出会带随机性。这个错误很隐蔽因为输出看起来差不多对但就是不稳定。一定要在导出前调用 eval()。第二个坑是 tokenizer 的特殊 token 处理。有些模型的 decoder_start_token 不是 pad_token如果搞错了生成的第一个 token 就是错的后面全乱。我建议导出前先把 tokenizer 的特殊 token id 都打印出来记在代码注释里。第三个坑是 KV 缓存的顺序。不同模型的 past_key_values 结构可能不一样有的是 (key, value) 交替有的是先所有 key 再所有 value。导出前一定要打印一下结构确认顺序。我因为这个搞反过一次解码结果完全不可用排查了大半天。第四个坑是最大长度设置。设太小会导致长句子翻译不完整设太大又浪费计算。我的经验是设成源句子长度的 2 到 2.5 倍比较合理再加一个硬上限比如 256。另外要确保循环里有 break 条件遇到 eos 就停不然会一直生成到最大长度。6. 进一步优化的几个方向6.1 批处理与动态 batch单句翻译效率低实际部署肯定要批处理。ONNX Runtime 支持动态 batch导出时把 batch 维度标成动态就行。但批处理有个问题不同句子长度不一样padding 到统一长度会浪费计算。解决办法是用 bucket 策略把长度相近的句子分到一组减少 padding。或者用 ONNX Runtime 的变长输入支持但那个配置比较复杂。我一般先用简单 padding等性能成为瓶颈再优化。6.2 模型量化与剪枝量化前面提过了动态量化最简单效果也不错。如果追求极致可以做静态量化需要校准数据集但精度和速度的平衡更好。剪枝是另一个方向把不重要的权重置零或者删掉。不过剪枝对翻译质量的影响比量化大需要仔细评估。我一般只在模型体积有硬性限制的时候才考虑剪枝。6.3 与其他推理引擎的对比除了 ONNX Runtime还有 TensorRT、OpenVINO 这些推理引擎。TensorRT 在 NVIDIA GPU 上性能最强但只支持特定硬件而且转换过程比 ONNX 麻烦。OpenVINO 在 Intel CPU 上表现好和 ONNX 的兼容性也不错。我的建议是先用 ONNX Runtime 跑通这是最通用的方案。如果确定部署硬件再针对性地换更专用的引擎。ONNX 作为中间格式可以很方便地转到其他引擎所以先把 ONNX 这步做好后面换引擎成本很低。6.4 部署时的工程化考虑实际部署不只是跑通推理还要考虑服务化、并发、监控这些。我一般用 FastAPI 或者类似的框架把推理包装成 HTTP 服务然后用多进程或者异步来处理并发请求。模型加载要放在服务启动时不要每次请求都加载。ONNX Runtime 的 InferenceSession 是线程安全的可以多个请求共享一个会话。但如果并发很高可以考虑每个进程一个会话用进程池来扩展。监控方面至少要记录每个请求的延迟、输入输出长度、有没有报错。这些数据能帮你发现性能瓶颈和异常情况。我还会定期抽样检查翻译质量确保模型没有退化。这个迁移方案我在几个项目里都跑过从最初的导出报错到最后的稳定部署大概花了两三天时间调试。最难的部分是 KV 缓存的导出和生成循环的编写但只要把这两块搞明白剩下的都是常规工程问题。如果你也在做类似的迁移建议先从一个小模型开始练手跑通整个流程后再换大模型这样能少走很多弯路。