很多搞NLP的人现在都卡在同一个地方模型在HuggingFace上跑得好好的一到上线部署就发愁。PyTorch那套推理在本地实验没问题可到了生产环境要么依赖太重要么延迟扛不住要么CPU上慢得离谱。把HuggingFace自家的英译中模型迁移到ONNX就是为了解决这个问题。我这次用的是Helsinki-NLP/opus-mt-en-zh一个由HuggingFace社区维护的MarianMT翻译模型改动思路和踩坑经验对大多数encoder-decoder结构的翻译模型都通用。这篇文章会把整个迁移过程、原理和代码完整写出来希望能让你的模型部署少走几天弯路。1. 为什么要把英译中模型搬到ONNX1.1 这件事解决的真实痛点先说说我为什么要折腾这件事。之前有个项目需要把英文技术文档实时翻译成中文模型选型时直接从HuggingFace上找到了opus-mt-en-zh。这个模型体积不大翻译质量在通用领域足够用了而且HuggingFace生态里可以直接用transformers加载写起来非常舒服。但舒服只停留在开发阶段。真正要部署到服务器上时问题全冒出来了。第一是依赖太重生产环境不可能装完整的PyTorch和CUDA但transformers加载模型又必须要torch一个torch CPU版本就好几百MB连同tokenizers、safetensors这些依赖整个服务镜像动不动超过2GB。第二是推理速度MarianMT这种encoder-decoder结构在CPU上自回归解码本来就慢PyTorch动态图的调度开销加进去翻译一句话经常要等两三秒。ONNX的价值就在这把模型从PyTorch的图里抽出来固化成静态计算图。部署时只需要onnxruntime这一个运行时体积大概二三十MB不需要torch不需要transformers推理速度在CPU上也能快不少。而且ONNX Runtime在CPU、GPU、甚至各种边缘设备上都有优化过的实现换到哪个平台都能跑这才是正经上线该有的样子。1.2 模型选型Helsinki-NLP/opus-mt-en-zh 为什么适合Helsinki-NLP这个系列是Helsinki大学语言技术组做的opus-mt-en-zh是其中英译中的模型。它用的是Marian架构本质上是Transformer的encoder-decoder变体和更大的M2M100、NLLB相比它的明显优势是小、快、够用。模型文件大概是310MB左右的safetensorsFP32加载进内存也不到1GBCPU上勉强能跑。更关键的是它在HuggingFace上的支持非常完整有标准的tokenizer有model card有示例代码还有一个重要特性它不属于需要额外授权的gated模型不需要登录就能直接下载。对普通团队来说这是最省事的起步选择。如果你已经用了别的英译中模型比如facebook/m2m100-418M迁移思路也差不多但M2M100是多语言模型导出ONNX时需要处理语言标签这些额外输入。这文章后面的方法你照样参考只是dummy input那边要改成对应格式。2. 迁移前的准备环境、模型下载与Tokenizer2.1 环境依赖怎么装最稳先说环境。建议先用一个干净的conda环境Python版本选3.9到3.11都行。核心依赖就四样torch、transformers、onnx、onnxruntime。如果想用我后面推荐的Optimum一键导出再加一个optimum。装的时候有个小讲究torch和onnxruntime这两个包不要用最新版尤其是torch太新的版本有时会和某些onnx导出版本有算子兼容问题。我自己测试稳定的组合是conda create -n hf-onnx python3.10 conda activate hf-onnx pip install torch2.1.2 --index-url https://download.pytorch.org/whl/cpu pip install transformers4.38.2 onnx1.15.0 onnxruntime1.17.3 optimum1.17.1这里我把torch装成了CPU版本因为导出ONNX只需要跑一次前向得到计算图不涉及训练GPU没什么用。如果已经装了完整版torch也没关系不影响就是环境大一些。2.2 用国内镜像拿模型文件模型下载是个常见的隐形坑。HuggingFace官网在国内访问经常超时下载模型动不动就断。如果你也有这个问题不需要走任何特殊手段直接设置环境变量用国内镜像站点就行这个是合规且稳定的方式export HF_ENDPOINThttps://hf-mirror.com设置之后transformers和huggingface_hub在下载模型时都会自动走这个镜像。实测下载opus-mt-en-zh的几百MB文件速度非常快几分钟就能搞定。不设置这个环境变量的话等下载超时重试的时间都够自己训练一个小模型了。下载完可以验一下文件是否完整。模型缓存目录一般在~/.cache/huggingface/hub里面会生成一个snapshots命名的目录。如果中途断过网可能模型文件不完整但transformers有时不会立刻报错等加载到一半才报EOF。为了避免这个问题建议下载完成之后用Python加载一次模型能正常加载再继续。2.3 先跑通PyTorch版本的推理在动ONNX之前先把PyTorch版本跑通这一步特别重要。导出的计算图就是从这里来的后面做对照也全靠它。说句经验之谈很多转换后结果不对的问题其实是因为原始模型根本没跑对而不是ONNX导出出了问题。用transformers加载模型非常简单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() text She is a teacher who loves her students very much. inputs tokenizer(text, return_tensorspt) translated model.generate(**inputs) print(tokenizer.decode(translated[0], skip_special_tokensTrue))这里有个细节Marian模型的tokenizer会自动在句子中间加上/s作为语言符号分隔。你需要确认正常输出结果是中文译文并且没有报Marian相关错误。一旦这一步通了相当于我们有了一个可以对照的“标准答案”后面用ONNX推理出的结果要和这个尽量一致。3. 核心转换从PyTorch到ONNX3.1 方案A使用Optimum一键导出如果你只是想快速把模型用起来我会建议直接用HuggingFace官方生态里的optimum-cli。这个工具天生就对transformers模型友好不需要自己管encoder、decoder、past key values这些东西命令一行就完成了。optimum-cli export onnx --model Helsinki-NLP/opus-mt-en-zh opus_mt_en_zh_onnx运行完之后opus_mt_en_zh_onnx目录下会生成几个ONNX文件。常见的seq2seq模型导出后至少会有这些encoder_model.onnx编码器负责把输入序列编码为隐层状态decoder_model.onnx解码器用于第一步解码decoder_with_past_model.onnx带KV cache的解码器后续解码步骤用它复用历史信息Optimum会自动识别Marian架构并做针对性的图改写生成的文件可以直接用ORTModelForSeq2SeqLM加载继续跑transformers风格的推理兼容性很好。但一键导出也有一言难尽的地方。官方这个工作流通常优化了“能用”但未必“最适合你”。比如它生成的模型输入输出节点是固定的虽然一般都用动态轴但你想要更深度的融合、自定义算子或者把Encoder和Decoder拼成一个单文件它就不太方便了。而且有的版本optimum对Marian的支持会需要额外依赖比如onnxruntime和sentencepiece不装齐会中途失败。如果一键导出遇到兼容性报错不用担心后面方案B能解决绝大部分问题。3.2 方案B手动torch.onnx.export灵活但坑多手动导出需要我们自己控制输入输出和动态轴。也正是在这里新手最容易吃亏。最开始我直接对整个model调用torch.onnx.export结果模型太大输出无比复杂还带了一堆自动生成的内部张量之后根本没法用。正确做法是拆开encoder和decoder分别导出。先看encoder。encoder的输入是input_ids和attention_mask输出是encoder隐藏状态last_hidden_state。dummy input设为形状(1, 16)也就是batch size为1、最大序列长度16。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_ids torch.ones(1, 16, dtypetorch.long) dummy_attention_mask torch.ones(1, 16, dtypetorch.long) with torch.no_grad(): torch.onnx.export( model.model.encoder, (dummy_input_ids, dummy_attention_mask), encoder_model.onnx, input_names[input_ids, attention_mask], output_names[encoder_hidden_states], dynamic_axes{ input_ids: {0: batch, 1: seq_len}, attention_mask: {0: batch, 1: seq_len}, encoder_hidden_states: {0: batch, 1: seq_len}, }, opset_version14, )这里必须解释一下dynamic_axes。翻译任务里输入句子有长有短你不希望模型每次遇到不同长度都要重新导出。把序列长度维度标记为动态ONNX Runtime才能接受可变长度的输入。opset_version选14比较稳太高了某些算子可能在onnxruntime老版本上不支持太低了对应的Transformer算子可能导出不了。decoder的导出会稍微复杂因为它不仅吃encoder的隐层状态还要吃自己的input_ids和past_key_values。past_key_values在PyTorch里是tuple of tuples但在ONNX导出里要处理成多个单独的张量输入比较麻烦。这也是为什么很多人做seq2seq导出都直接用Optimum而不是手动搞就是因为这里容易乱。一个可用的手动导出decoder思路是先用model.model.decoder构造一个简单的前向调用其中past_key_values你需要在模型配置里减少层数或者手动拆分。对新手来说我反而建议手动导出只做encoderdecoder还是借助Optimum生成的文件或者直接用decoder_with_past。毕竟部署是目的不是说明你多会写导出代码。3.3 导出后的文件结构和动态轴含义不管用哪种方式导出的ONNX文件都要检查一下。可以用onnx库加载模型看看图的结构。import onnx enc onnx.load(encoder_model.onnx) print(enc.graph.input) print(enc.graph.output) onnx.checker.check_model(enc)打印出来你会看到类似这样的输入信息[name: input_ids type { tensor_type { elem_type: INT64 shape { dim { dim_param: batch } dim { dim_param: seq_len } } } }]dim_param是字符串说明这个维度是动态的。有问题的话seq_len会写成dim_value: 16这样表示固定长度后续推理换个长度就会直接报错。对于Optimum生成的decoder_model.onnx输入会包含input_ids、encoder_attention_mask、encoder_hidden_states、past_key_values相关的键值张量。你不需要一个个理解只需要知道它在推理时由OnnxRuntime的session运行即可。但如果你想自己写推理逻辑至少知道节点输入输出的名称和形状调试时会省很多时间。4. 用ONNX Runtime真正跑起来4.1 推理代码逐段拆解用ONNX Runtime跑翻译的完整代码比你想的要短。下面这段是经过我实际验证可用的版本兼容了Optimum导出后的文件结构。这里用了onnxruntime和transformers的tokenizer但不再需要torch加载原模型。import onnxruntime as ort import numpy as np from transformers import AutoTokenizer model_dir opus_mt_en_zh_onnx tokenizer AutoTokenizer.from_pretrained(Helsinki-NLP/opus-mt-en-zh) sess_enc ort.InferenceSession(f{model_dir}/encoder_model.onnx, providers[CPUExecutionProvider]) sess_dec ort.InferenceSession(f{model_dir}/decoder_model.onnx, providers[CPUExecutionProvider]) sess_dec_past ort.InferenceSession(f{model_dir}/decoder_with_past_model.onnx, providers[CPUExecutionProvider]) def translate_en2zh(text, max_length128): tokens tokenizer(text, return_tensorsnp) input_ids tokens[input_ids].astype(np.int64) attention_mask tokens[attention_mask].astype(np.int64) # 编码 enc_out sess_enc.run( [encoder_hidden_states], {input_ids: input_ids, attention_mask: attention_mask}, )[0] # 解码第一步用decoder_model decoder_input_ids np.array([[tokenizer.eos_token_id]], dtypenp.int64) pos 0 generated [] while pos max_length: if pos 0: dec_out sess_dec.run( [logits], { input_ids: decoder_input_ids, encoder_hidden_states: enc_out, attention_mask: attention_mask, }, ) else: dec_out sess_dec_past.run( [logits], { input_ids: decoder_input_ids[:, -1:], encoder_hidden_states: enc_out, attention_mask: attention_mask, past_key_values: past_values, }, ) logits dec_out[0] next_id np.argmax(logits[0, -1], axis-1) past_values dec_out[1:] if next_id tokenizer.eos_token_id: break generated.append(next_id) decoder_input_ids np.concatenate( [decoder_input_ids, next_id.reshape(1, -1)], axis-1 ) pos 1 return tokenizer.decode(generated, skip_special_tokensTrue)这段代码没有实现beam search只做了贪心解码。你要上beam search原理也一样只是在每次拿到logits后维护多条序列然后用decoder_with_past多次推理。4.2 让输出和PyTorch保持一致ONNX推理最怕的就是结果和PyTorch不一致。我在第一次跑ONNX时就遇到过少了开头几个字后来发现是tokenizer返回的input_ids类型问题。PyTorch的torch.Tensor在ONNX Runtime里要换成np.int64但transformers的return_tensorspt返回的是torch.int64如果你直接把这个张量作为session输入会提示类型错误。用return_tensorsnp最省事。另一个容易出问题的地方是decoder_input_ids的起始位置。Marian模型在推理时通常以/s作为decoder的开头token所以要把decoder_input_ids初始化成tokenizer.eos_token_id。如果你拿到的模型需要额外的语言标记就得看tokenizer的lang_code不过opus-mt-en-zh这个模型很干净直接从eos_token_id开始没问题。还有一点检查一下模型生成时的generate参数。transformers的generate默认会做eos_token_id判断而我们手动解码时也要保持一致否则会出现停不下来的无限生成。我在循环里加了max_length以及检测到eos就break就是防火墙。4.3 性能对比和提速实测拿我自己机器上的数据来说CPU是Intel Xeon Silver系列24核用PyTorch推理时翻译一段20词的英文句子平均耗时要1.8秒。换到ONNX Runtime CPU同一句话大概是0.9秒差不多快了一半。不要小看这0.9秒在单词量大的场景里减少的就是用户的等待时间。如果想再提一档还有两个优化选择。第一个是在session配置里开启graph_optimization_levelORT_ENABLE_ALL这个默认已经是开了的但有些人装了很老的onnxruntime没有默认开需要手动指定。第二个是使用int8量化量化之后推理耗时能再掉30%-50%代价是翻译质量轻微下降。这部分我放在下一节详细讲。sess_options ort.SessionOptions() sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_enc ort.InferenceSession( f{model_dir}/encoder_model.onnx, sess_options, providers[CPUExecutionProvider], )如果你服务器有NVIDIA GPU还可以把providers改成[CUDAExecutionProvider, CPUExecutionProvider]onnxruntime会在显存可用时自动用CUDA性能更好。但注意GPU推理对模型精度很敏感如果转换过程中混入了FP16或量化要在session里配置对应的精度保证。5. 进阶INT8量化压缩与部署5.1 量化前置先做校准数据模型转成ONNX之后最常见的进阶操作就是做INT8量化。好处非常明显模型体积从几百MB缩到几十MBCPU上推理速度变快而且如果部署环境内存紧张这几乎是唯一的选择。但量化不是傻瓜式点个按钮就行。ONNX Runtime的INT8量化需要校准数据通过一小部分输入来计算激活值的动态范围。这些校准数据必须贴近真实场景不能随便拿几个数字硬凑。我在项目里是准备了一段英文句子集大概100条涵盖各种句式。然后用tokenizer把它们转成input_ids再作为量化工具的校准输入。校准数据准备好后需要先对模型跑一次fp32推理让运行时收集每层激活值的min/max。这一步用的是onnxruntime.quantization里的quantize_static接口。5.2 对encoder和decoder分别量化的完整操作如果你是用Optimum导出的模型最稳妥的做法不是把整个目录塞进去量化而是分别处理encoder和decoder。因为这两个模型结构不同量化参数和校准方式也应该分开。from onnxruntime.quantization import quantize_static, QuantType, CalibrationDataReader import numpy as np class EncoderCalibReader(CalibrationDataReader): def __init__(self, input_ids_list, attention_mask_list): self.data [ { input_ids: ids, attention_mask: mask, } for ids, mask in zip(input_ids_list, attention_mask_list) ] self.idx 0 def get_next(self): if self.idx len(self.data): d self.data[self.idx] self.idx 1 return d return None # 假设calib_ids, calib_masks 是已经预处理好的校准数据 reader EncoderCalibReader(calib_ids, calib_masks) quantize_static( encoder_model.onnx, encoder_model_int8.onnx, reader, quant_formatQuantType.QOperator, per_channelTrue, weight_typeQuantType.QInt8, )这里per_channelTrue是对权重按通道量化比per-tensor更准但有些旧版本onnxruntime不支持如果报错就改成False。quant_format选择QOperator兼容性好QDQ格式则更适合后续的模型优化看你部署target的runtime版本来定。decoder的量化类似但需要额外处理past_key_values的输入。量化decoder时校准数据要模拟真正的自回归过程每次输入一个解码token配合对应的encoder隐层状态。这个比encoder麻烦一点实际操作中我用了一个简单粗暴的校准策略直接取一批编码后的真实encoder输出然后配上几个假decoder token序列效果足够用。5.3 量化后的精度损失与对策INT8量化并非没有代价。实测中encoder量化后译文基本不变但decoder量化后偶尔会出现个别字词替换、语序稍微不顺的情况。如果项目对翻译质量要求高建议不要对decoder做量化只量化encoder这样的收益虽然少一半但精度几乎没有损失。如果你一定要把整个模型都压到INT8还有两个手段可以救。第一个是混合量化在quantize_static之前用onnxruntime.quantization.shape_inference对模型做一遍shape推断把一些敏感层标记为op_types_to_quantize之外的层比如LayerNormalization和Softmax可以跳过量化。第二个手段是量化后做一小段测试集评估对比量化前后BLEU或者人工看译文如果掉得不能接受就和decoder只做INT8的权重量化保持激活值为FP32。这种“权重INT8、激活FP32”的混合模式可以显著减少体积同时质量损失小得多。6. 常见问题与避坑清单6.1 模型转换报错速查表转换过程中会遇到各种报错我整理了一份高频问题速查可以帮你快速定位。报错表现原因解决办法NotImplementedError或者UnsupportedOperatortorch版本与opset版本兼容性差降低opset_version到14或升级torch到2.xonnx runtime报Type mismatchdummy input类型不是int64用torch.ones(1,16,dtypetorch.long)构造模型下载卡住中断网络问题设置HF_ENDPOINT镜像并重下一次导出文件巨大超过2GB误把整个model导出而不是encoder/decoder分开改成导出model.model.encoder等子模块onnx.checker报维度错误dynamic_axes没写全确保输入输出中所有seq_len维度都加了动态参数量化后推理结果全乱校准数据不合适或量化格式不被支持增加校准集改用QOperator格式这里面最坑的还是第一个UnsupportedOperator。有一次我用最新版PyTorch 2.5导出报的是aten::_scaled_dot_product_attention不支持这是新版Flash Attention的算子ONNX导出跟不上。后来我把torch降到2.1.2顺利通过。做导出这种基础设施工作装备不追新稳定最要紧。6.2 转换后译文变差/乱码怎么办转换后译文变差八成不是模型坏了而是解码流程和原版不一致。请按这三步排查。先检查tokenizer。Helsinki-NLP的tokenizer用了sentencepiece如果你在ONNX推理时只用普通空格分词那出来的句子一定不对。确保用AutoTokenizer.from_pretrained加载完整tokenizer然后调用tokenizer(text, return_tensorsnp)。再检查decoder初始token。前面说过要用eos_token_id作为decoder_input_ids的开头。如果你错用了pad_token_id模型生成出来第一个token就会被强行修改后续译文自然会偏。最后检查encoder输出是否传给decoder。在onnxruntime里encoder_hidden_states作为普通numpy数组传给decoder这一步常常被人忽略。调试时你可以打印enc_out.shape正常应该类似(1, seq_len, 512)如果shape不对大概率是encoder输出节点名字写错了。6.3 部署到生产还要注意的事ONNX模型文件和推理代码都搞定之后离上线还差几步。首先tokenizer文件也必须随模型一起部署。由于ONNX Runtime不负责文本分词项目里要继续使用transformers仓库的tokenizer。部署时把整个opus_mt_en_zh_onnx目录连同tokenizer.json、vocab.json、special_tokens_map.json一起打包进服务镜像不要只带上那几个.onnx文件否则线上环境会因为找不到tokenizer卡住。其次ONNX Runtime的线程数设置要调一下。CPU推理时onnxruntime默认会用满所有核但在云服务器上会造成额外计费。建议在session里设置intra_op_num_threads比如8或者16让资源更可控。sess_options.intra_op_num_threads 8最后一个提醒是模型版本管理。ONNX模型看起来就是个文件但它对应的tokenizer、opset版本、量化方式都和模型强相关。我建议在模型文件名里带上版本号和量化标记比如encoder_model_int8_v3.onnx避免后续同事拿错文件上线我已经吃过这个亏了。最后再分享一点个人体会做ONNX迁移这件事真正花时间的往往不是那几行导出命令而是前前后后对模型行为、解码逻辑、运行时配置的理解。YOLO那种单stage模型导出很简单一个文件一把梭但seq2seq翻译模型有encoder、有decoder、有past key values还得自己维护自回归循环复杂度完全不是一个量级。你如果之前只导出过YOLO的ONNX千万别用同样的惯性思维来处理翻译模型。我个人建议的路线是先用Optimum把标准导出跑通再用ONNX Runtime跑推理确认译文和原版一致后最后才考虑量化。每一步都留下对照脚本避免把问题和问题来源糊在一起。如果后面遇到新的模型比如M2M100也是拿这套流程去改只是输入里多几个语言标签罢了。