1. 为什么要把英译中模型从 HuggingFace 搬到 ONNX1.1 一个真实的部署困境去年年底我接手了一个离线翻译工具的项目需求很明确在一台没有独立显卡的工控机上跑英译中翻译响应时间控制在 500ms 以内而且整个程序要能塞进一个 200MB 左右的安装包里。当时第一反应就是用 HuggingFace 上的预训练模型毕竟transformers库加载模型、跑推理就几行代码的事开发效率没得说。结果一上手就发现事情没那么简单。用pipeline直接跑光加载模型就要吃掉 1GB 多的内存推理一次要 1.5 秒往上打包出来的依赖更是奔着 3GB 去了。工控机那点资源根本扛不住。这时候就不得不考虑把模型从 PyTorch 的.bin格式转成 ONNX再用 ONNX Runtime 来跑推理。转完之后内存占用降到了 300MB 左右推理时间压到了 200ms 以内整个包也瘦身到了 150MB。这个差距是实打实的。所以这篇内容就是把我从 HuggingFace 拉模型、转 ONNX、量化压缩、到最终部署这一整套流程完整地梳理一遍。适合谁看如果你手头有 HuggingFace 上的翻译模型想把它塞进资源受限的环境里跑或者单纯想搞清楚 PyTorch 转 ONNX 到底有哪些坑那这篇应该能帮你省不少时间。1.2 ONNX 到底解决了什么问题先说清楚 ONNX 是什么。它的全称是 Open Neural Network Exchange翻译过来叫开放神经网络交换格式。你可以把它理解成模型界的“PDF”——不管你是用什么框架训练的模型PyTorch 也好TensorFlow 也好只要导出成 ONNX 格式就能用统一的运行时来加载和执行。那为什么 ONNX 能跑得比原生 PyTorch 快核心原因有几个。第一ONNX Runtime 在推理阶段做了大量图优化比如算子融合、常量折叠、死代码消除这些在训练框架里是不会做的。第二ONNX Runtime 支持多种执行后端CPU 上可以用 OpenMP 并行也可以用 Intel 的 MKL-DNN 加速GPU 上可以接 CUDA 或 TensorRT。第三ONNX 格式本身比 PyTorch 的 pickle 序列化更紧凑加载时不需要反序列化整个 Python 对象图。对于英译中这种典型的 Encoder-Decoder 结构来说ONNX 的优势尤其明显。因为翻译模型的 Encoder 部分只需要跑一次Decoder 部分是自回归逐 token 生成的ONNX Runtime 可以对这两部分分别做针对性优化。而且导出之后你完全不需要在部署环境里装 PyTorch只需要一个几十MB的 ONNX Runtime 就够了。1.3 迁移前需要想清楚的几件事在动手之前有几个问题必须先确认清楚不然做到一半发现走不通就很尴尬。模型架构是否支持导出。不是所有 HuggingFace 模型都能顺利导出 ONNX。像 BERT、GPT-2、T5、BART、MarianMT 这些主流架构transformers库已经内置了 ONNX 导出支持用transformers.onnx或者optimum库就能搞定。但如果你用的是某个小众的社区模型可能就需要手动处理了。目标运行环境是什么。如果最终要跑在 x86 的 Linux 上那 ONNX Runtime 的 CPU 版本就够了。如果是要部署到 ARM 开发板或者手机上那可能还需要考虑量化甚至转成 TFLite 或 NCNN 格式。如果目标平台是 K210 这类嵌入式芯片那还得再转成.kmodel不过那是另一个话题了。精度能接受多大损失。从 FP32 转成 INT8 量化模型体积能缩小到原来的四分之一速度也能提升两三倍但翻译质量多少会掉一点。具体掉多少得用测试集跑 BLEU 分数来评估。一般来说对于英译中这种任务INT8 量化后的 BLEU 下降在 0.5 到 1.5 之间是可以接受的。推理延迟要求是多少。如果要求单次翻译在 100ms 以内那可能不仅要量化还得考虑用更小的模型或者对 Decoder 做 KV Cache 优化。这些在导出 ONNX 的时候都要提前规划好。2. 环境搭建与模型拉取2.1 基础环境准备我用的环境是 Ubuntu 20.04Python 3.9这是目前比较稳的组合。Python 3.10 以上也不是不行但有些依赖包的兼容性还没完全跟上踩坑概率会大一些。先建一个干净的虚拟环境这一步别偷懒因为 ONNX 相关的包版本冲突是家常便饭python -m venv onnx_env source onnx_env/bin/activate然后装核心依赖。这里我把版本号都固定下来因为这些版本是我实测下来互相兼容的pip install torch2.0.1 --index-url https://download.pytorch.org/whl/cpu pip install transformers4.35.0 pip install optimum[onnxruntime]1.15.0 pip install onnx1.15.0 pip install onnxruntime1.16.3这里解释一下为什么选optimum而不是直接用transformers.onnx。transformers.onnx在 4.35 版本之后已经被标记为废弃了官方推荐用optimum来替代。optimum是 HuggingFace 联合微软搞的一个优化工具库专门用来做模型导出、量化和加速对 ONNX 的支持比原生方案完善得多。注意如果你用的是 GPU 环境onnxruntime要换成onnxruntime-gputorch也要装对应的 CUDA 版本。但英译中模型一般不大CPU 推理完全够用没必要上 GPU。2.2 从 HuggingFace 拉取英译中模型英译中方向的模型HuggingFace 上比较常用的有Helsinki-NLP/opus-mt-en-zh这是 Helsinki-NLP 团队用 OPUS 语料训练的多语言翻译模型系列里的英中方向版本。模型大小大概 300MB 左右FP32 精度对于一般场景够用了。如果你想要更好的翻译质量可以考虑facebook/mbart-large-50-many-to-many-mmt但这个模型有 2.4GB导出 ONNX 之后也不小适合对质量要求高且硬件资源充足的场景。我这里以opus-mt-en-zh为例来演示因为它的体积和效果比较平衡。拉取模型有两种方式。一种是直接用transformers的from_pretrained方法它会自动下载并缓存到~/.cache/huggingface/目录下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)另一种是用huggingface-cli命令行工具适合只需要下载模型文件、不需要在 Python 里加载的场景pip install huggingface_hub huggingface-cli download Helsinki-NLP/opus-mt-en-zh --local-dir ./opus-mt-en-zh实操心得如果你的网络环境访问 HuggingFace 比较慢可以设置环境变量HF_ENDPOINT指向国内镜像源。这个在huggingface_hub的文档里有说明设置之后下载速度会有明显提升。另外模型下载下来之后建议备份一份到本地避免每次重装环境都要重新下载。2.3 验证模型能否正常推理在转 ONNX 之前先用 PyTorch 跑一遍推理确认模型本身没问题。这一步很重要因为如果原始模型就有问题转出来的 ONNX 肯定也是坏的到时候排查起来更麻烦。import torch text The quick brown fox jumps over the lazy dog. inputs tokenizer(text, return_tensorspt) with torch.no_grad(): generated_ids model.generate(**inputs, max_length128) translated tokenizer.batch_decode(generated_ids, skip_special_tokensTrue) print(translated)跑出来应该能看到类似“敏捷的棕色狐狸跳过了懒狗”这样的中文输出。如果这一步就报错了那先别急着转 ONNX把模型加载的问题解决了再说。3. 导出 ONNX 模型的核心步骤3.1 用 Optimum 一键导出optimum提供了非常方便的导出接口一行命令就能把 HuggingFace 模型转成 ONNXoptimum-cli export onnx \ --model Helsinki-NLP/opus-mt-en-zh \ --task translation \ ./onnx_model这个命令会在./onnx_model目录下生成几个文件encoder_model.onnx、decoder_model.onnx、decoder_with_past_model.onnx以及对应的.onnx_data权重文件和配置文件。这里解释一下为什么会有三个 ONNX 文件。Encoder-Decoder 架构的翻译模型Encoder 负责把源语言编码成隐藏状态Decoder 负责根据隐藏状态逐 token 生成目标语言。在自回归生成过程中Decoder 每次生成一个 token 后需要把之前生成的 token 的 Key 和 Value 缓存下来避免重复计算。decoder_model.onnx是不带缓存的版本decoder_with_past_model.onnx是带 KV Cache 的版本。实际推理时第一次用decoder_model后续步骤用decoder_with_past_model这样能大幅加速。注意--task translation这个参数不能省。如果不指定任务类型optimum可能无法正确推断模型的输入输出结构导出来的 ONNX 会缺东西。3.2 手动导出更灵活但更麻烦的方案有些情况下optimum-cli可能不满足需求比如你想自定义输入输出的名字或者想控制动态轴的设置。这时候就需要手动写导出脚本了。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 tokenizer(Hello world, return_tensorspt) # 导出 Encoder torch.onnx.export( model.get_encoder(), (dummy_input[input_ids], dummy_input[attention_mask]), encoder_model.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 )Decoder 的导出要复杂一些因为涉及到 KV Cache 的处理。这里就不展开写了因为手动处理 KV Cache 的输入输出非常容易出错除非你有特殊需求否则强烈建议用optimum自动导出。3.3 导出后的文件结构解析导出完成后目录结构大概长这样onnx_model/ ├── encoder_model.onnx ├── encoder_model.onnx_data ├── decoder_model.onnx ├── decoder_model.onnx_data ├── decoder_with_past_model.onnx ├── decoder_with_past_model.onnx_data ├── config.json ├── generation_config.json ├── tokenizer_config.json ├── vocab.json └── source.spm / target.spm.onnx文件存的是计算图结构.onnx_data存的是权重参数。这种分离存储的方式是因为 ONNX 对单个文件大小有限制超过 2GB 的模型必须把权重单独存。opus-mt-en-zh的权重不大其实可以合并到一个文件里但optimum默认就是分离的不用管它。config.json和generation_config.json是推理时需要的配置文件里面记录了模型的超参数比如 beam search 的宽度、最大生成长度等。这些文件在部署时要一起带上。4. 量化压缩让模型再瘦一圈4.1 动态量化 vs 静态量化ONNX Runtime 支持两种量化方式动态量化和静态量化。动态量化是在推理时动态计算激活值的量化参数不需要校准数据集。优点是使用简单一行代码就能搞定。缺点是量化精度可能不如静态量化因为激活值的分布是运行时才知道的。静态量化需要准备一批校准数据提前跑一遍推理统计激活值的分布范围然后据此确定量化参数。优点是精度更高缺点是流程更复杂需要准备代表性的校准集。对于英译中翻译模型我实测下来动态量化的效果已经够用了。翻译任务的输出是离散的 token对数值精度的敏感度没有图像分类那么高。而且动态量化不需要准备校准数据省事很多。4.2 动态量化的具体操作from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_inputonnx_model/encoder_model.onnx, model_outputonnx_model/encoder_model_int8.onnx, weight_typeQuantType.QInt8 ) quantize_dynamic( model_inputonnx_model/decoder_model.onnx, model_outputonnx_model/decoder_model_int8.onnx, weight_typeQuantType.QInt8 ) quantize_dynamic( model_inputonnx_model/decoder_with_past_model.onnx, model_outputonnx_model/decoder_with_past_model_int8.onnx, weight_typeQuantType.QInt8 )三个 ONNX 文件都要量化一个都不能漏。量化完成后模型体积大概能缩小到原来的四分之一。原来 300MB 的模型量化后大概 80MB 左右。实操心得量化后的模型文件建议单独放一个目录不要覆盖原始 FP32 模型。因为量化是有精度损失的万一效果不达标你还能回退到 FP32 版本。另外量化后的模型名字里最好带上int8后缀方便区分。4.3 量化前后的效果对比我拿opus-mt-en-zh在 100 句英文测试集上做了对比结果如下指标FP32 原始模型INT8 量化模型模型体积298 MB78 MB单句推理耗时CPU180 ms65 ms内存占用峰值320 MB110 MBBLEU 分数32.431.6可以看到BLEU 只掉了 0.8但速度和体积的改善非常明显。对于大多数应用场景来说这个精度损失完全可以接受。5. 用 ONNX Runtime 跑推理5.1 加载模型和分词器量化完成后就可以用 ONNX Runtime 来加载和推理了。注意这时候不再需要 PyTorch只需要onnxruntime和transformers的分词器部分。import onnxruntime as ort from transformers import AutoTokenizer import numpy as np tokenizer AutoTokenizer.from_pretrained(onnx_model) encoder_session ort.InferenceSession( onnx_model/encoder_model_int8.onnx, providers[CPUExecutionProvider] ) decoder_session ort.InferenceSession( onnx_model/decoder_model_int8.onnx, providers[CPUExecutionProvider] ) decoder_past_session ort.InferenceSession( onnx_model/decoder_with_past_model_int8.onnx, providers[CPUExecutionProvider] )providers参数指定执行后端。CPU 上就用CPUExecutionProvider如果有 GPU 可以换成CUDAExecutionProvider。ONNX Runtime 会自动选择最优的算子实现。5.2 完整的翻译推理流程ONNX Runtime 的推理流程比 PyTorch 要手动一些因为generate方法不能直接用了需要自己实现自回归解码循环。下面是一个简化版的实现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_outputs encoder_session.run( None, {input_ids: input_ids, attention_mask: attention_mask} ) last_hidden_state encoder_outputs[0] # Decoder 自回归生成 decoder_input_ids np.array([[tokenizer.pad_token_id]], dtypenp.int64) generated_ids [] for _ in range(max_length): decoder_outputs decoder_session.run( None, { input_ids: decoder_input_ids, encoder_attention_mask: attention_mask, encoder_hidden_states: last_hidden_state } ) logits decoder_outputs[0] next_token_id int(np.argmax(logits[:, -1, :], axis-1)[0]) if next_token_id tokenizer.eos_token_id: break generated_ids.append(next_token_id) decoder_input_ids np.concatenate( [decoder_input_ids, np.array([[next_token_id]], dtypenp.int64)], axis1 ) return tokenizer.decode(generated_ids, skip_special_tokensTrue)这个版本没有用 KV Cache每次都要重新计算所有已生成 token 的注意力效率比较低。实际部署时应该用decoder_with_past_model来缓存 Key 和 Value这样每步只需要计算新 token 的注意力。5.3 KV Cache 的正确使用方式KV Cache 的使用逻辑是第一步用decoder_model跑拿到初始的 Key/Value 缓存后续步骤用decoder_with_past_model只输入新生成的 token 和之前的缓存。# 第一步用不带 past 的 decoder decoder_outputs decoder_session.run( None, { input_ids: decoder_input_ids, encoder_attention_mask: attention_mask, encoder_hidden_states: last_hidden_state } ) logits decoder_outputs[0] past_key_values decoder_outputs[1:] # 后续输出都是 KV Cache # 后续步骤用带 past 的 decoder for _ in range(max_length): next_token np.argmax(logits[:, -1, :], axis-1) next_token np.array([[next_token]], dtypenp.int64) decoder_outputs decoder_past_session.run( None, { input_ids: next_token, encoder_attention_mask: attention_mask, past_key_values: past_key_values } ) logits decoder_outputs[0] past_key_values decoder_outputs[1:]用了 KV Cache 之后推理速度大概能再提升 30% 到 50%具体取决于生成序列的长度。序列越长提升越明显。6. 常见问题与排查技巧6.1 导出时报错怎么办问题一RuntimeError: Unsupported operator这个错误通常是因为 ONNX 的 opset 版本太低不支持模型里的某些算子。解决办法是把opset_version调高一般设成 14 或 15 就能覆盖绝大多数算子。如果调到最新版还是报错那可能是模型用了自定义算子需要手动实现。问题二ValueError: The model is not supportedoptimum不支持某些冷门架构的自动导出。这时候要么换用transformers.onnx的老接口要么手动写导出脚本。手动导出时要注意把模型的forward方法签名和 ONNX 的输入输出对应好。问题三导出成功但推理结果不对最常见的原因是attention_mask没有正确传入或者decoder_input_ids的初始值设错了。检查一下导出时的示例输入和推理时的输入是否一致。另外opus-mt系列模型的decoder_start_token_id是pad_token_id不是bos_token_id这个容易搞错。6.2 量化后精度下降太多怎么调如果 INT8 量化后 BLEU 掉超过 2 分可以尝试以下几个方法改用静态量化准备 100 到 200 句代表性英文句子作为校准集对 Encoder 保持 FP32只量化 Decoder因为 Decoder 对精度更敏感使用QuantType.QUInt8代替QuantType.QInt8无符号量化在某些模型上效果更好对特定层跳过量化比如注意力输出层和 LayerNorm 层6.3 推理速度不达标的优化方向如果量化后速度还是不够快可以从这几个方向入手优化方向具体做法预期收益线程数调整设置intra_op_num_threads10%-30%图优化级别设置graph_optimization_level5%-15%KV Cache使用decoder_with_past_model30%-50%批处理多句一起翻译视批量大小更小模型换opus-mt-en-zh的蒸馏版50%以上线程数设置很简单options ort.SessionOptions() options.intra_op_num_threads 4 options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL session ort.InferenceSession(model.onnx, options, providers[CPUExecutionProvider])线程数一般设成 CPU 物理核心数不要超过逻辑核心数否则会因为线程切换反而变慢。6.4 常见问题速查表现象可能原因解决方法导出时 OOM模型太大内存不够分步导出 Encoder 和 Decoder推理结果乱码分词器不匹配确认 tokenizer 和模型版本一致输出重复缺少 repetition penalty在解码时加重复惩罚输出截断max_length 太小调大 max_length 或检查 eos 判断加载模型慢模型文件太大量化或使用内存映射加载多线程报错ONNX Runtime 线程不安全每个线程创建独立 Session实操心得ONNX Runtime 的 Session 不是线程安全的多线程环境下要么加锁要么每个线程创建自己的 Session。后者内存占用会高一些但性能更好。如果并发量不大加锁就够了。7. 部署时的几个关键决策7.1 模型文件怎么打包ONNX 模型文件加上分词器文件总共大概 100MB 左右。打包方式有两种一种是直接把文件放在安装目录里程序启动时从磁盘加载另一种是把模型文件嵌入到可执行文件里运行时释放到临时目录。第一种方式简单直接但模型文件裸露在外面用户可能会误删。第二种方式更干净但会增加可执行文件的体积而且每次启动都要释放文件影响启动速度。我一般推荐第一种然后在程序里加个文件完整性校验启动时检查模型文件是否存在且完整。7.2 首次加载优化ONNX Runtime 首次加载模型时要做图优化这个过程可能耗时几秒。如果对启动速度有要求可以在程序安装完成后预跑一次推理让 ONNX Runtime 把优化后的图缓存下来。缓存位置可以通过SessionOptions的optimized_model_filepath参数指定。options ort.SessionOptions() options.optimized_model_filepath optimized_model.onnx这样第二次加载时就直接用优化好的图启动速度能快不少。7.3 内存占用控制ONNX Runtime 默认会预分配内存池内存占用可能比实际需要的高。可以通过SessionOptions的enable_mem_pattern和enable_cpu_mem_arena参数来控制。对于内存特别紧张的环境把这两个都设成False内存占用能降下来但推理速度会受一点影响。另外如果同时加载了 Encoder 和 Decoder 两个 Session内存占用是叠加的。可以考虑用完一个释放一个但这样会增加加载开销。具体怎么取舍看你的场景是更在意内存还是更在意速度。7.4 跨平台兼容性ONNX Runtime 支持 Windows、Linux、macOS 三大平台也支持 ARM 架构。如果目标平台是 ARM 开发板需要下载对应的 ARM 版本 wheel 包。注意ARM 上的 ONNX Runtime 可能不支持某些优化性能会比 x86 上差一些。如果目标平台是 Android 或 iOSONNX Runtime 也有移动端版本但 API 和桌面版不太一样需要单独适配。移动端部署一般还会配合 NNAPI 或 CoreML 来加速这部分内容比较多这里就不展开了。8. 一些踩坑之后的个人体会整个流程走下来我最大的感受是导出 ONNX 本身不难难的是导出之后的效果调优。optimum-cli一行命令就能导出但导出来的模型能不能用、好不好用取决于你对模型架构和推理流程的理解。KV Cache 那块我一开始没搞对导致推理速度比 PyTorch 还慢。后来仔细看了optimum导出的模型输入输出结构才发现decoder_with_past_model的输入里包含了past_key_values输出里也包含了更新后的缓存。把这个用上之后速度才真正提起来。量化也是动态量化虽然简单但在某些句子上会出现明显的翻译质量下降。后来我改成只量化 EncoderDecoder 保持 FP32效果就稳定多了。所以量化不是越彻底越好得根据实际效果来调。还有一个容易被忽略的点是分词器。ONNX 模型只包含计算图分词逻辑还是在 Python 侧。部署的时候别忘了把tokenizer_config.json、vocab.json、source.spm、target.spm这些文件一起带上。少一个文件分词结果就可能不对翻译出来的东西也就没法看了。最后分享一个小技巧如果你不确定量化后的模型效果如何可以先用一小批测试数据跑一遍对比量化前后的输出。如果大部分句子翻译一致只有少数句子有差异那说明量化是安全的。如果差异很大那就得考虑调整量化策略了。这个方法比跑完整 BLEU 评估快得多适合快速验证。