1. 为什么要把英译中模型从 HuggingFace 搬到 ONNX1.1 一个真实的需求场景去年帮一个做跨境电商的朋友处理商品详情页的本地化问题他们的场景很典型每天有几千条英文商品描述需要翻译成中文早期直接调用在线翻译接口量小的时候没什么感觉量一上来账单就很难看了。后来他们改成自己部署模型用 HuggingFace 上的开源英译中模型比如Helsinki-NLP/opus-mt-en-zh这类效果够用成本也可控。但新的问题很快出现了。原始模型是 PyTorch 格式推理时依赖完整的 PyTorch 运行时一个进程吃内存不少启动还慢。他们想把这套东西塞进一个边缘设备做离线翻译或者至少在一台普通服务器上多开几个实例扛并发PyTorch 那套就显得太重了。这时候把模型迁移到 ONNX 就成了一个很自然的选择。ONNX 全称 Open Neural Network Exchange你可以把它理解成深度学习模型的“通用中间格式”。PyTorch、TensorFlow 训练出来的模型都能转成 ONNX然后交给 ONNX Runtime 去跑。ONNX Runtime 是个专门做推理的引擎不带训练相关的包袱所以体积小、启动快、内存占用低还能针对不同硬件做图优化。对于英译中这种 encoder-decoder 结构的翻译模型来说迁移到 ONNX 之后推理速度通常能有明显提升尤其是在 CPU 上。这篇文章我就把整个迁移过程拆开讲清楚从环境准备、模型导出、分词器处理到推理验证、量化加速再到实际部署时踩过的坑。适合已经会用 HuggingFace 加载模型、但没怎么碰过 ONNX 的同学也适合想把翻译模型塞进生产环境的工程师。1.2 迁移前必须搞清楚的几个概念在动手之前有几个概念必须先理清楚不然很容易在中间某个环节卡住。第一个是 encoder-decoder 结构。英译中模型基本都是 seq2seq 架构包含一个 encoder 负责理解英文输入一个 decoder 负责逐词生成中文。这跟 BERT 那种只有 encoder 的模型不一样导出 ONNX 的时候要处理的东西更多尤其是 decoder 的自回归生成过程。第二个是 KV Cache。翻译模型生成中文时是一个词一个词往外蹦的每生成一个新词decoder 都要重新计算一遍之前所有词的注意力。如果不做缓存重复计算量会非常大。KV Cache 就是把之前算过的 key 和 value 存下来复用这是 decoder 推理提速的关键。导出 ONNX 时要不要带 KV Cache直接决定了推理代码的复杂度。第三个是分词器。HuggingFace 的模型都配了对应的 tokenizer英文侧和中文侧的分词方式可能完全不同。ONNX 只负责神经网络那部分计算分词、编码、解码这些前后处理还是得靠 Python 代码或者别的工具来做。很多人迁移完发现结果不对问题往往出在分词器上而不是模型本身。第四个是动态轴。翻译的输入句子长度是不固定的导出 ONNX 时必须把序列长度这类维度标记成动态的否则模型只能处理固定长度的输入那就没法用了。把这四个概念想明白后面的操作就顺了。2. 环境准备与工具选型2.1 依赖安装的取舍环境这块其实不复杂核心就几个包。我一般用 conda 建一个干净的虚拟环境避免和系统里的其他包打架。conda create -n hf2onnx python3.10 conda activate hf2onnx pip install torch transformers onnx onnxruntime版本上有个经验transformers和torch的版本要匹配不然导出时可能报一些莫名其妙的错。我实测下来transformers4.30配合torch2.0比较稳。onnx和onnxruntime建议都装最新版因为 ONNX 的算子集在持续更新老版本可能不支持某些新算子。如果你打算做量化还要额外装onnxruntime-tools或者用optimum这个库。optimum是 HuggingFace 官方出的导出工具封装得比较好但我个人更倾向于手动导出因为可控性强出问题好排查。新手可以先用optimum跑通流程再回头研究手动导出的细节。提示不要在生产环境直接pip install一堆包建议用requirements.txt锁版本。ONNX 相关的包版本兼容性比较敏感今天能跑的代码明天换个版本可能就报错。2.2 模型下载的现实问题HuggingFace 的模型在国内下载有时候会比较慢这是很多人第一步就卡住的地方。有几个办法可以缓解一是用huggingface-cli配合镜像站点二是提前把模型文件下载到本地再加载。pip install huggingface_hub huggingface-cli download Helsinki-NLP/opus-mt-en-zh --local-dir ./opus-mt-en-zh下载完之后加载模型时直接指定本地路径就行from transformers import AutoTokenizer, AutoModelForSeq2SeqLM model_path ./opus-mt-en-zh tokenizer AutoTokenizer.from_pretrained(model_path) model AutoModelForSeq2SeqLM.from_pretrained(model_path)这样做的另一个好处是导出 ONNX 的过程可以完全离线进行不依赖网络。生产环境部署时也是同样的思路把模型文件打包进镜像或者挂载到容器里。2.3 选哪个模型做迁移HuggingFace 上的英译中模型不少常见的有Helsinki-NLP/opus-mt-en-zh、facebook/mbart-large-50-many-to-many-mmt等。选哪个主要看你的场景。opus-mt-en-zh体积小大概 300MB 左右推理快适合对延迟敏感、翻译质量要求不是极致的场景。mbart-large-50质量更好但模型大很多导出 ONNX 之后文件也大推理慢。我这次拿opus-mt-en-zh做演示因为它结构典型迁移过程有代表性而且大部分人的实际需求用这个就够了。选模型的时候还要注意一点有些模型是 multilingual 的tokenizer 里有一大堆语言标记导出时这些细节都要照顾到。单语言的英译中模型相对简单适合第一次做迁移练手。3. 核心细节解析导出 ONNX 的关键环节3.1 导出脚本的整体结构导出 ONNX 的核心代码其实不长但每一行都有讲究。先看一个能跑通的最小版本import torch from transformers import AutoTokenizer, AutoModelForSeq2SeqLM model_path ./opus-mt-en-zh tokenizer AutoTokenizer.from_pretrained(model_path) model AutoModelForSeq2SeqLM.from_pretrained(model_path) model.eval() # 构造一个示例输入 dummy_text Hello, how are you today? inputs tokenizer(dummy_text, return_tensorspt) # 导出 encoder torch.onnx.export( model.get_encoder(), (inputs[input_ids], inputs[attention_mask]), 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, )这里有几个关键点。model.eval()必须调用否则 dropout 之类的层会处于训练模式导出的模型行为不对。dynamic_axes把 batch 和 sequence 两个维度标成动态的这样模型才能处理任意长度、任意批量的输入。opset_version建议用 14 或更高低版本对某些算子支持不好。encoder 的导出相对简单因为它就是一次前向计算输入输出都很明确。真正麻烦的是 decoder。3.2 Decoder 导出的两种方案Decoder 的导出有两种思路各有优劣我分别说一下。方案一不带 KV Cache 的完整 decoder。每次生成一个词都把完整的输入序列重新喂给 decoder 算一遍。这种方案导出简单推理代码也简单但速度慢因为重复计算太多。适合对速度要求不高、只想先跑通的场景。方案二带 KV Cache 的 decoder。把之前算过的 key、value 缓存起来每次只算新词。速度快很多但导出和推理代码都复杂要手动管理 cache 的输入输出。生产环境基本都用这个方案。我建议先用方案一跑通整个流程确认翻译结果正确再切换到方案二做优化。直接上方案二容易在 cache 的维度处理上翻车。带 KV Cache 的导出核心是要把 past_key_values 作为输入输出暴露出来。transformers里 decoder 的 forward 方法支持use_cacheTrue导出时要把这些 cache 张量都列进input_names和output_names。层数乘以 2key 和 value就是 cache 张量的数量opus-mt-en-zh有 6 层 decoder所以有 12 组 cache。3.3 分词器的处理ONNX 模型只吃数字不吃文本。所以推理时英文输入要先经过 tokenizer 编码成 input_ids模型输出 logits 之后还要经过 tokenizer 解码成中文。这部分逻辑必须和导出时保持一致否则结果会错乱。有个容易忽略的点opus-mt-en-zh的 tokenizer 在编码时可能会自动加一些特殊 token比如cmn_Hans这种语言标记。导出时用的示例输入和推理时用的输入处理方式必须完全一样。我一般会把 tokenizer 的配置也保存下来推理时重新加载确保一致。tokenizer.save_pretrained(./onnx_tokenizer)推理时tokenizer AutoTokenizer.from_pretrained(./onnx_tokenizer)这样即使换了环境分词行为也不会变。3.4 动态轴设置的细节动态轴设置错了是导出后推理报错的高频原因。除了 batch 和 sequence有些模型还有额外的维度需要标记。比如 attention_mask 在某些实现里形状是[batch, 1, 1, sequence]那就要把对应的维度也标成动态。一个实用的排查方法导出后用onnxruntime加载模型打印输入输出的形状信息看看哪些维度是固定的。如果发现某个本该动态的维度是固定值就回去改dynamic_axes。import onnxruntime as ort sess ort.InferenceSession(encoder.onnx) for inp in sess.get_inputs(): print(inp.name, inp.shape) for out in sess.get_outputs(): print(out.name, out.shape)形状里出现字符串比如batch、sequence的就是动态维度出现具体数字的就是固定维度。对照你的需求检查一遍能省掉很多调试时间。4. 实操过程从导出到推理验证4.1 完整导出流程把前面的碎片拼起来一个完整的导出脚本大概是这样。我把它分成 encoder 和 decoder 两部分decoder 先用不带 cache 的简单版本。import torch from transformers import AutoTokenizer, AutoModelForSeq2SeqLM model_path ./opus-mt-en-zh tokenizer AutoTokenizer.from_pretrained(model_path) model AutoModelForSeq2SeqLM.from_pretrained(model_path) model.eval() dummy_text Hello, how are you today? 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), encoder.onnx, input_names[input_ids, attention_mask], output_names[encoder_hidden_states], dynamic_axes{ input_ids: {0: batch, 1: sequence}, attention_mask: {0: batch, 1: sequence}, encoder_hidden_states: {0: batch, 1: sequence}, }, opset_version14, ) # 导出 decoder不带 cache decoder model.get_decoder() decoder_input_ids torch.tensor([[tokenizer.pad_token_id]]) encoder_hidden_states model.get_encoder()(input_ids, attention_mask)[0] torch.onnx.export( decoder, (decoder_input_ids, encoder_hidden_states), decoder.onnx, input_names[decoder_input_ids, encoder_hidden_states], output_names[logits], dynamic_axes{ decoder_input_ids: {0: batch, 1: decoder_sequence}, encoder_hidden_states: {0: batch, 1: sequence}, logits: {0: batch, 1: decoder_sequence}, }, opset_version14, )导出完成后目录下会多出encoder.onnx和decoder.onnx两个文件。文件大小加起来应该和原始 PyTorch 模型差不多如果差太多可能是导出时漏了什么。4.2 用 ONNX Runtime 做推理导出只是第一步能不能跑出正确结果才是关键。下面是一段完整的推理代码实现了贪心解码import numpy as np import onnxruntime as ort from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(./onnx_tokenizer) encoder_sess ort.InferenceSession(encoder.onnx) decoder_sess ort.InferenceSession(decoder.onnx) def translate(text, max_length50): inputs tokenizer(text, return_tensorsnp) input_ids inputs[input_ids].astype(np.int64) attention_mask inputs[attention_mask].astype(np.int64) encoder_hidden_states encoder_sess.run( None, {input_ids: input_ids, attention_mask: attention_mask} )[0] # 从 decoder 的起始 token 开始 decoder_input_ids np.array([[tokenizer.pad_token_id]], dtypenp.int64) generated [] for _ in range(max_length): logits decoder_sess.run( None, { decoder_input_ids: decoder_input_ids, encoder_hidden_states: encoder_hidden_states, } )[0] next_token int(np.argmax(logits[0, -1, :])) if next_token tokenizer.eos_token_id: break generated.append(next_token) decoder_input_ids np.concatenate( [decoder_input_ids, np.array([[next_token]], dtypenp.int64)], axis1 ) return tokenizer.decode(generated, skip_special_tokensTrue) print(translate(Hello, how are you today?))跑通之后应该能输出类似“你好你今天怎么样”的结果。如果输出是乱码或者空字符串八成是分词器或者起始 token 的问题。4.3 参数选择与性能对比我在这台机器上做了个简单的对比测试输入是一段 20 个词的英文句子输出限制 50 个 tokenCPU 环境。方案单句耗时内存占用模型文件大小PyTorch 原始模型约 420ms约 1.2GB约 310MBONNX 不带 cache约 380ms约 600MB约 300MBONNX 带 cache约 150ms约 650MB约 320MB可以看到不带 cache 的 ONNX 相比 PyTorch 提升有限主要省在内存和启动时间上。带 cache 之后速度提升就很明显了接近三倍。所以如果你的场景对延迟敏感一定要上 KV Cache。内存占用这块ONNX Runtime 比 PyTorch 轻不少这对多实例部署很关键。同样一台机器PyTorch 可能只能开两三个实例ONNX 能开五六个。4.4 量化加速的实操如果还想进一步压榨性能可以做 int8 量化。ONNX 的量化分动态量化和静态量化动态量化最简单不需要校准数据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, )量化之后模型文件能小一半左右推理速度在支持 int8 指令集的 CPU 上还能再快一些。但要注意量化会带来一定的精度损失翻译质量可能略有下降。我实测下来opus-mt-en-zh量化后质量下降不明显日常使用基本感觉不出来但如果你的场景对翻译质量要求极高建议先做一轮对比测试再决定。注意量化后的模型在某些老 CPU 上可能反而变慢因为缺少 int8 指令集支持。部署前一定要在目标硬件上实测。5. 常见问题与排查技巧实录5.1 导出阶段的高频报错导出阶段最容易遇到的是算子不支持的问题。transformers里有些操作在导出时会被拆成 ONNX 不认识的算子报错信息通常是Unsupported operator或者Exporting operator xxx failed。解决办法有几个一是提高opset_version新版本算子集支持更多操作二是用torch.onnx.export的trainingtorch.onnx.TrainingMode.EVAL参数确保导出的是推理图三是实在不行就改模型代码把不支持的操作替换成等价的、ONNX 支持的操作。还有一个常见报错是RuntimeError: Expected all tensors to be on the same device。这通常是模型在 GPU 上、但示例输入在 CPU 上导致的。导出时统一用 CPU 就行ONNX 导出不依赖 GPU。5.2 推理结果不对的排查思路推理结果不对排查顺序建议是先查分词器再查输入形状最后查模型逻辑。分词器问题最隐蔽。比如英文输入编码后input_ids 里可能包含一些特殊 token如果推理时没处理模型看到的就是错的输入。我的做法是把 PyTorch 版本的推理结果和 ONNX 版本的推理结果逐 token 对比看从哪一步开始出现差异。输入形状问题也好查。ONNX Runtime 对输入的形状和类型要求很严格int64和int32混用都会报错。用sess.get_inputs()打印出期望的类型然后确保喂进去的数据类型一致。模型逻辑问题相对少见但如果前面都排查过了还是不对就要检查 decoder 的起始 token 是不是正确。不同的模型起始 token 不一样有的是pad_token_id有的是decoder_start_token_id用错了整个生成过程就偏了。5.3 常见问题速查表问题现象可能原因解决办法导出报 Unsupported operatoropset 版本低提高到 14 或以上推理报 shape mismatch动态轴没设对检查 dynamic_axes 配置输出乱码分词器不一致保存并复用导出时的 tokenizer输出为空起始 token 错误用 decoder_start_token_id速度没提升没带 KV Cache改用带 cache 的 decoder量化后变慢CPU 不支持 int8换硬件或放弃量化内存持续增长cache 没释放每次推理后重置 cache5.4 几个踩过的坑坑一batch 推理时 padding 处理。单句推理没问题一批句子一起推理就出错多半是 padding 没对齐。ONNX 模型要求同一批的输入长度一致需要手动 pad 到相同长度同时 attention_mask 要正确标记哪些是 padding。坑二decoder 的 cache 维度。带 cache 的 decoderpast_key_values 的形状是[batch, num_heads, past_seq_len, head_dim]第一次推理时 past_seq_len 是 0。这个 0 长度的张量在 ONNX 里处理起来比较麻烦有些版本会报错。解决办法是第一次推理时用一个长度为 1 的 dummy cache或者干脆第一轮不带 cache。坑三中文解码的编码问题。有些 tokenizer 解码出来的中文是带空格的比如“你 好”需要额外处理。这个不是 ONNX 的问题是 tokenizer 的特性用skip_special_tokensTrue和clean_up_tokenization_spacesTrue能缓解。坑四多线程推理的线程安全问题。ONNX Runtime 的 session 是线程安全的但如果你在多个线程里共享同一个 session 并手动管理 cache就可能出问题。生产环境建议每个线程独立管理自己的 cache 状态。6. 部署上线的几点经验6.1 服务化封装模型导出完最终还是要包成一个服务。最简单的做法是用 FastAPI 包一层 HTTP 接口from fastapi import FastAPI from pydantic import BaseModel import onnxruntime as ort from transformers import AutoTokenizer app FastAPI() tokenizer AutoTokenizer.from_pretrained(./onnx_tokenizer) encoder_sess ort.InferenceSession(encoder.onnx) decoder_sess ort.InferenceSession(decoder.onnx) class TranslateRequest(BaseModel): text: str app.post(/translate) def translate(req: TranslateRequest): result do_translate(req.text) return {result: result}ONNX Runtime 的 session 在服务启动时初始化一次就行不要每次请求都重新加载那样开销太大。session 本身是线程安全的多个请求可以共享。6.2 性能调优的几个开关ONNX Runtime 提供了一些性能相关的配置值得调一调。intra_op_num_threads控制单个算子内部的并行线程数inter_op_num_threads控制算子之间的并行。CPU 环境下把这两个值设成物理核心数通常效果最好。options ort.SessionOptions() options.intra_op_num_threads 4 options.inter_op_num_threads 2 options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess ort.InferenceSession(encoder.onnx, options)graph_optimization_level设成ORT_ENABLE_ALL会启用所有图优化包括算子融合、常量折叠等对推理速度有帮助。这些参数没有万能值建议在目标硬件上做一轮压测找到最适合的配置。6.3 版本管理与回滚ONNX 模型文件建议和代码一起做版本管理。每次导出新模型记录下用的 transformers 版本、opset 版本、导出脚本的 commit这样出问题能快速定位。生产环境部署时新模型先灰度一小部分流量确认没问题再全量。我个人习惯在模型文件名里带上版本号和日期比如encoder_v2_20240115.onnx避免覆盖旧文件。回滚的时候直接切回旧文件就行不用重新导出。6.4 后续还能怎么优化如果这套流程跑通了还有几个方向可以继续深挖。一是尝试不同的量化策略静态量化配合校准数据通常比动态量化效果更好。二是探索把模型转成其他推理引擎支持的格式比如针对特定硬件的优化格式进一步压榨性能。三是把 encoder 和 decoder 合并成一个 ONNX 图减少两次 session 调用的开销不过这样灵活性会差一些。另外如果你的场景是固定的一批翻译任务可以考虑做批处理优化把多个句子攒成一批一起推理吞吐量能提升不少。这个在 ONNX 里实现起来比 PyTorch 更直接因为 ONNX Runtime 对批处理的调度做得比较好。我在实际部署中发现真正影响线上表现的往往不是模型本身而是前后处理的效率。分词、padding、解码这些 Python 层面的操作如果写得不够高效可能比模型推理还慢。所以优化的时候别只盯着模型把整条链路都测一遍找到真正的瓶颈再动手。