ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

HuggingFace英译中模型迁移ONNX实战:从PyTorch到CPU高效部署

HuggingFace英译中模型迁移ONNX实战:从PyTorch到CPU高效部署 我是在一次内部工具开发里产生那个非常“上头”的需求HuggingFace 上的英译中开源模型质量很能打但把它放到生产环境的 CPU 服务器上却让我体验了一把“能跑、但没法用”的落差。项目名很直白就是“将 HuggingFace 自己的英译中模型迁移到 ONNX”。这个活儿做完之后模型推理速度、部署体积、工程化程度都上升了一大截。这篇文章我打算把手摸过的所有流程、踩过的坑、最终落地的方案完整捋一遍包括基于 Helsinki-NLP/opus-mt-en-zh 这个 MarianMT 模型的迁移全过程。不管你是第一次了解 ONNX还是已经在做 NLP 模型部署应该都能从这里找到可以直接抄作业的步骤。1. 为什么非迁不可英译中模型的部署痛点1.1 直接跑 PyTorch CPU 推理的糟糕体验先说一个我最想吐槽的场景模型质量明明没问题但每次启动一个翻译接口光是加载 torch、transformers 这一层依赖内存就已经吃掉了接近 1GB。然后在 8 核 CPU 上翻译一句 30 个 token 左右的英文PyTorch FP32 推理耗时大概在 800ms 到 1.2s 之间这还不算 Python 侧 generate 循环的调度开销。生产环境里如果同时来几个请求服务器基本就处于“看起来活着但响应时间完全不可控”的状态。另外就是版本管理问题。transformers 的版本、torch 的版本、GPU 驱动和 CUDA 的版本任何一个变动都可能导致推理结果变化甚至是直接跑不起来。而 ONNX 模型是纯计算图加载它只需要一个 onnxruntime 运行时不依赖整个深度学习框架这就把“环境兼容”这件事从根上简化了。很多团队会把模型放到专门的推理服务里但小团队没有那么多资源去维护一个独立推理集群。把这些模型迁移到 ONNX 之后用一个轻量级 Python 服务甚至 C 服务就能把翻译能力嵌入现有系统这才是最“省心”的部署方式。1.2 ONNX 到底解决了什么问题ONNX 本身是“开放神经网络交换格式”它做的事情就是把训练好的模型冻结成一张静态计算图。图里面节点是算子边是张量数据。你不需要管模型当初是用 PyTorch、TensorFlow 还是其他框架训练的只要它被顺利转换成了 ONNX就可以用统一的运行时去执行。对比 PyTorch 的动态图机制ONNX 最大的优势集中在三块运行时轻量onnxruntime 是 C 实现Python 端只是封装安装体积小、启动速度快、内存占用低。CPU 推理优化充分它会把算子级别的优化、图级别的优化都做掉比如算子融合、常数折叠、内存复用很多场景下比直接用 PyTorch 推理快 30% 到 50%。便于量化与剪枝ONNX 生态有专门的结构化量化和动态量化工具几行代码就能把 FP32 模型压到 INT8体积和速度都能进一步优化。我当时比较过同一句话的翻译延迟从 PyTorch 到 ONNX FP32 再到 ONNX INT8整个下降曲线非常直观后面我会放一份对比数据。1.3 英译中模型选型MarianMT 是最合适的起点HuggingFace 上有不少英译中模型常见的有 Helsinki-NLP 系的 MarianMT、Facebook 系的 M2M100 和 NLLB-200。NLLB-200 翻译质量和多语言覆盖确实猛但模型规模太大CPU 部署基本不现实。M2M100 也是类似情况性价比不适合小项目。MarianMT 的opus-mt-en-zh是专门的英语到中文方向的模型模型体积在 300MB 左右由 6 层 encoder 和 6 层 decoder 组成。它的翻译质量在通用文本上能打七八十分日常工具完全够用而且模型结构相对简单非常适合作为“迁移到 ONNX”的练手对象。选型的时候我建议遵循一个原则先选你能跑得动的模型再追求高质量模型。ONNX 迁移本身并不改变模型参数和结构只改变推理执行的方式。如果模型太大CPU 上跑不动转成 ONNX 也只是杯水车薪。2. 环境准备与模型结构拆解2.1 依赖安装比想象中简单这个项目需要的依赖并不多核心就是 transformers、onnx、onnxruntime再加一个 optimum 作为辅助工具。我自己用的是 Python 3.10PyTorch 2.x。pip install transformers torch onnx onnxruntime optimum如果希望用 optimum 的命令行工具导出模型还需要确认optimum安装了 exporters 相关模块一般装完之后会自动带上。版本需要注意一下onnxruntime 建议用较新版本算子兼容性会好很多尤其是导出时用到 Transformer 内部一些新算子的时候老版本容易报错。2.2 下载模型与网络问题的处理HuggingFace 的模型下载本来是直接一行代码就完事但由于服务端在国内访问经常会出现超时、断连、下载到一半失败的情况我这边的经验是不要在交互环境里反复“裸奔”下载而是先用huggingface_hub的缓存机制把模型完整拉下来再离线加载。pip install huggingface_hub python -c from huggingface_hub import snapshot_download; snapshot_download(Helsinki-NLP/opus-mt-en-zh)如果网络环境不稳下载可以多试几次断点续传一般能帮你省不少时间。模型文件不大耐心等一会儿就好。下载完成后模型会缓存到本地目录后续代码里即使不联网也能通过from_pretrained正常加载。2.3 MarianMT 的结构与 KV Cache 概念理解 MarianMT 的结构是整个迁移过程的转折点。它属于经典的 Encoder-Decoder 架构Encoder 负责把源语言句子编码成一组上下文向量Decoder 负责逐词生成目标语言文本。我最初犯过一个认知错误以为这种模型可以像 BERT 一样一次性把所有 token 的预测结果给出来。实际上生成是逐词的每一步都依赖之前已经生成的 token这就是所谓的“自回归”生成。自回归生成里有一个尤其关键的概念就是 KV Cache。Decoder 在预测第 n 个词的时候其实只需要关注当前这一步的输入和前面已经算好的 key、value 缓存不需要重算整个序列。PyTorch 的generate内部会自动管理这个缓存但到了 ONNX 手工推理时你就必须自己维护 past_key_values也就是把 model 输出的所有 key/value 张量保存下来在下一次推理时作为输入再喂回去。我把这个逻辑理解了之后后面写推理代码就顺多了。如果你也想把别的 Seq2Seq 模型迁移到 ONNX这个结构认知是通用的。3. Pytorch 转 ONNX 的两种实操路径3.1 最快方案一行命令用 optimum 导出如果你不想管太多底层细节直接上 optimum 的命令行工具就行。第一次跑通的时候我甚至有点意外整个过程远比自己想象中顺滑。optimum-cli export onnx --model Helsinki-NLP/opus-mt-en-zh --opset 13 onnx_model/执行完成之后onnx_model目录下会出现这几个文件encoder_model.onnxdecoder_model.onnxdecoder_with_past_model.onnx配套的 config.json、tokenizer 文件等这三个 ONNX 文件分别对应编码器、解码器无缓存版和解码器带缓存版。decoder_with_past_model.onnx就是专门用于自回归推理的高性能版本每一步只需要输入上一步生成的 token 以及上一步的 KV Cache推理效率明显更优。这个方式的优点是省心可靠optimum 已经帮你处理了 MarianDecoder 里各种复杂输出的包装问题。缺点是你容易变成“只会敲命令的搬运工”一旦模型结构稍微特殊一点或者推理代码写得不顺仍然需要退回底层去理解计算图。3.2 手动导出方案Encoder 先上手手动导出可以帮助你看清每个环节的原理。我的实验过程是先做 encoder再做 decoder。MarianMT 的 encoder 输入是input_ids和attention_mask输出是last_hidden_state。直接取model.model.encoder包装成子模块导出import torch from transformers import MarianMTModel model MarianMTModel.from_pretrained(Helsinki-NLP/opus-mt-en-zh) encoder model.model.encoder encoder.eval() dummy_input torch.zeros((1, 32), dtypetorch.long) dummy_mask torch.ones((1, 32), dtypetorch.long) torch.onnx.export( encoder, (dummy_input, dummy_mask), encoder_model.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_version13, )这里最需要留意的是dynamic_axes。如果不声明 seq_len 是动态维度导出的模型会固定死在 seq_len32推理时一旦句子长度变了就会报错。opset_version13是当前兼容性和算子支持都非常均衡的版本不要为了“更新”就去追太高的版本部分 ONNX Runtime 的老版本可能不支持最新 opset 里的新算子。3.3 手动导出方案Decoder 的包装与 KV CacheDecoder 比 Encoder 麻烦得多主要原因是它的输入除了input_ids之外还要接收来自 encoder 的上下文向量并且内部需要处理 KV Cache。直接拿model.model.decoder裸导出通常会失败因为 Transformer 源码里的叠加输出结构tuple of tuple不是 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_valuesNone, use_cacheTrue, ): outputs self.decoder( input_idsinput_ids, encoder_hidden_statesencoder_hidden_states, encoder_attention_maskencoder_attention_mask, past_key_valuespast_key_values, use_cacheuse_cache, ) logits outputs[0] past outputs[1] keys [] values [] for layer_past in past: keys.append(layer_past[0]) values.append(layer_past[1]) return logits, *keys, *values直接把 flatten 后的 key、value 作为图输出这样导出的 ONNX 图输入输出都规规矩矩后面用 onnxruntime 推理时只需要按顺序把这些张量重新组织回 past 字典即可。这里要特别说明手动导出时past_key_valuesNone的初始步和past_key_values非空的后续步最好分别导出成两个模型一个给首轮用一个给迭代轮用。我导出来的首版 decoder 模型只适合首轮推理循环里继续复用会导致结果错误。手动导出 Decoder 的代码只是展示核心思路实际落地时optimum已经帮我们封装好了这些细节所以如果你赶时间直接用 3.1 的产物就好。如果你要迁移的模型比较特殊optimum 不支持那这个包装思路就是你的突破口。3.4 动态轴与 opset 选择的门道关于动态轴我单独提一下。encoder、decoder 的input_ids和attention_mask的序列维度必须设成动态英文句子长度和中文句子长度不一定成比例你没办法在部署时预测最大长度。实际操作中我会把 batch 维度也保留为动态这样同一个模型既能单条请求也能批量请求。opset 的选择上13 是稳妥牌。它支持大部分 Transformer 内部算子又不至于引入太多新特性导致 runtime 版本要求过高。之前我在一个旧环境里用过 opset 17 导出的模型结果那边 onnxruntime 版本偏老加载直接失败那个教训记忆犹新。4. ONNX Runtime 推理与验证4.1 先打印输入输出名别盲写推理代码模型导出之后很多人的第一反应是照着网上的瞎写喂数据写完一跑直接 shape mismatch。我的建议是先把三个 ONNX 模型的输入输出名字全打印出来核对一遍再动手。import onnxruntime as ort for name in [encoder_model.onnx, decoder_model.onnx, decoder_with_past_model.onnx]: session ort.InferenceSession(fonnx_model/{name}) inputs [(i.name, i.shape) for i in session.get_inputs()] outputs [(o.name, o.shape) for o in session.get_outputs()] print(name) print(inputs:, inputs) print(outputs:, outputs) print(- * 30)这一步看起来很基础却帮我避开了至少两处命名不一致的问题。不同模型、不同 transformers 版本导出的 ONNX 输入名可能有细微差别比如encoder_hidden_states有时叫encoder_outputs。拿到真实名字之后再构造 feeds出错的概率要低很多。4.2 完整推理代码带 KV Cache 优化下面给出我最终落地的推理代码。这个版本直接使用原生的 onnxruntime没有依赖 optimum 的包装类核心逻辑就是“首轮走 decoder_model后续轮次走 decoder_with_past_model”。import numpy as np import onnxruntime as ort from transformers import MarianTokenizer tokenizer MarianTokenizer.from_pretrained(Helsinki-NLP/opus-mt-en-zh) sess_opt ort.SessionOptions() sess_opt.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_opt.intra_op_num_threads 4 enc_sess ort.InferenceSession(onnx_model/encoder_model.onnx, sess_optionssess_opt) dec_sess ort.InferenceSession(onnx_model/decoder_model.onnx, sess_optionssess_opt) dec_past_sess ort.InferenceSession(onnx_model/decoder_with_past_model.onnx, sess_optionssess_opt) def translate(text: str, max_len: int 128): enc tokenizer(text, return_tensorsnp) input_ids enc[input_ids] attention_mask enc[attention_mask] enc_out enc_sess.run([last_hidden_state], { input_ids: input_ids, attention_mask: attention_mask, })[0] eos_id tokenizer.eos_token_id pad_id tokenizer.pad_token_id decoder_input np.array([[pad_id]], dtypenp.int64) past None generated [] for _ in range(max_len): if past is None: feeds { input_ids: decoder_input, encoder_hidden_states: enc_out, encoder_attention_mask: attention_mask, } session dec_sess else: feeds { input_ids: decoder_input, encoder_hidden_states: enc_out, encoder_attention_mask: attention_mask, } for i, tensor in enumerate(past): feeds[fpast_key_values.{i // 2}.{key if i % 2 0 else value}] tensor session dec_past_sess results session.run(None, feeds) logits results[0] if len(results) 1: past results[1:] next_id int(np.argmax(logits[0, -1, :])) if next_id eos_id: break generated.append(next_id) decoder_input np.array([[next_id]], dtypenp.int64) return tokenizer.decode(generated, skip_special_tokensTrue)这个代码我在多个模型上都跑通过核心细节在于 decoder 每一步的attention_mask始终使用的是源语言的attention_mask而不是目标语言的。很多新手这里会弄混导致中文翻译结果前半段正常、后半段崩溃。如果你的模型导出时past_key_values 的命名和我的不一致记得根据 4.1 的打印结果调整 feeds 里的 key 名。4.3 与 PyTorch 输出对比验证迁移正确性迁移最怕的是“模型能跑但输出和原来不一样”。我踩过这个坑之后总结出一个经验任何模型迁移首先要做输出一致性验证。验证方式很简单用原版 PyTorch 模型和 ONNX 模型分别对同一组英文句子做 greedy decoding然后对比最终翻译结果。正常情况下FP32 的 ONNX 输出应该和 PyTorch 输出完全一样因为计算图里的算子几乎不会改变浮点运算的语义。from transformers import MarianMTModel pt_model MarianMTModel.from_pretrained(Helsinki-NLP/opus-mt-en-zh) pt_model.eval() sentences [ Hello, how are you today?, The ONNX runtime is very fast on CPU., I enjoyed reading this technical blog., ] for s in sentences: pt_out pt_model.generate(**tokenizer(s, return_tensorspt)) pt_text tokenizer.decode(pt_out[0], skip_special_tokensTrue) ort_text translate(s) print(PT :, pt_text) print(ORT:, ort_text) print()如果发现最终输出不一致优先检查两个点第一模型是否处于 eval 模式第二dropout 是否被意外开启。ONNX 模型里没有 dropout 随机性如果 PyTorch 端忘开 eval两者必然对不上。4.4 想用 Beam Search 怎么办上面的代码是 Greedy Search每次只取概率最高的 token。实际翻译场景里Beam Search 能明显提升句子通顺度。我一般在句子偏长、语法复杂的时候用 beam size 3 的效果比 greedy 好不少。实现思路不算复杂维护一个 beam 状态列表每个 beam 有自己的 generated token 序列和自己的 past_key_values每步对当前所有 beam 的 logits 做 top-k 选择把候选按累计概率重新排序裁掉概率低的 beam 之后继续。不过如果你想尽快上线也可以用 optimum 的ORTModelForSeq2SeqLM它内置了完整的 generate 逻辑支持 beam search、长度惩罚等参数。但我个人建议手工原生推理值得写一遍因为只有亲手管理过 past_key_values你才能真正理解 self-attention 的自回归过程。5. 性能优化INT8 量化与线程调参5.1 动态 INT8 量化实操模型导出完只是开始性能不优化等于白干。我首推的动态量化是 ONNX Runtime 里最成熟的一条路把权重从 FP32 量化到 INT8 或 UINT8推理时动态反量化整体精度损失在可接受范围内。量化只需要几行代码from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic(onnx_model/encoder_model.onnx, onnx_model/encoder_model_int8.onnx, weight_typeQuantType.QInt8) quantize_dynamic(onnx_model/decoder_model.onnx, onnx_model/decoder_model_int8.onnx, weight_typeQuantType.QInt8) quantize_dynamic(onnx_model/decoder_with_past_model.onnx, onnx_model/decoder_with_past_model_int8.onnx, weight_typeQuantType.QInt8)量化后会生成新的 ONNX 文件推理时把 session 指向 INT8 文件即可。注意这里的动态量化不是训练后完全静态量化它主要压缩权重激活值仍然是动态计算。胜在操作简单不需要校准数据集。5.2 Runtime 调参线程与图优化ONNX Runtime 的性能上下限差距很大关键在于 session 配置。我最常用的两个参数是intra_op_num_threads和graph_optimization_level。intra_op_num_threads控制单次算子里内部并行的线程数。在 CPU 服务器上盲目开满线程并不一定更快因为线程切换本身也有开销。我实际测下来8 核机器上设成 4 到 6 线程效果最好。graph_optimization_level推荐设为ORT_ENABLE_ALL它会启用算符融合和内存复用等各种优化。so ort.SessionOptions() so.intra_op_num_threads 4 so.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL5.3 优化前后数据对比以下是我在一台 8 核 Intel CPU 服务器上翻译一句约 20 个英文单词、最终生成约 40 个中文字符的实测数据方案模型文件体积单句平均耗时相对 PyTorchPyTorch FP32约 310MB约 850ms1.0xONNX FP32约 310MB约 460ms0.54xONNX INT8约 85MB约 320ms0.38x这个数据是给大家一个感知就算不做量化单是迁移到 ONNXCPU 推理速度就有接近翻倍的提升。量化和线程调优之后还能再快 30% 左右而且模型文件变小不少部署时加载模型的耗时也随之下降。5.4 量化注意质量与速度的平衡动态 INT8 量化不是无痛的。我在一个较长的口语化翻译任务里发现量化后的译文偶尔会丢掉一两个虚词通顺度略受影响。这是因为 INT8 压缩了权重精度导致 logits 的细微差异在 top-1 选择时可能造成跳跃式改变。解决方案有两个第一上线之前用小批量测试集对比量化前后翻译质量如果偏差异常大就只量化 encoderdecoder 保持 FP32第二尝试QuantType.QUInt8在某些层上的表现会比 QInt8 稳一点。实际工程中“部分量化”是性价比最高的方案encoder 层量化收益明显且质量影响小decoder 对精度更敏感可以先不动。6. 常见问题与避坑指南6.1 导出时报 Unsupported Operator这个是最常见的坑。旧版 transformers 的 MarianDecoder 内部某些实现会触发 PyTorch 独有的算子比如ATen系列。遇到这种问题第一件事是升级 transformers、onnx、onnxruntime 三件套然后设置opset_version13再导出。如果升级后仍然报错可以把相关算子的影响范围缩小。我的经验是绝大多数情况下升级到最新稳定版都能解决。不要一上来就怀疑模型结构有问题因为这通常不是模型的问题而是导出链路版本兼容的问题。6.2 动态轴没设置好导致 shape 报错如果推理时输入句子长度超过导出时固定的长度onnxruntime 会报诸如“Tried to access input position xxx but batch size is 1”之类非常绕的错。这时回到导出脚本检查dynamic_axes是否把input_ids的 seq_len 维度设成了动态。很多人会漏掉attention_mask的动态轴只设置input_ids导致后面 encoder 输出和 mask 形状对不上。我的建议是凡是涉及“输入 token 序列”和“输出 hidden state 序列”的维度全部设置为动态宁可多设也不能少设。6.3 tokenizer 不一致导致乱码ONNX 模型是不管 tokenizer 的英文和中文字符的切词规则必须由模型配套的 tokenizer 处理。我试过用别的 tokenizer 加载同一个模型结果中文输出全是 [UNK]。这件事提醒我保存 ONNX 模型时一定要把原始 tokenizer 的 vocab.json、special_tokens_map.json 文件一并拷贝到部署目录推理端也要显式加载同一个 tokenizer。6.4 past 名字顺序与维度对不上使用 decoder_with_past 推理时最容易翻车的就是把 past 张量按错误顺序塞进去。MarianMT 是 6 层 decoder过去缓存有 12 个张量分别按 layer0.key、layer0.value、layer1.key、layer1.value 交替排列。我建议不要在代码里硬编码名字而是像 4.1 那样打印输入输出名用循环动态构造 feeds。6.5 session 并发与生命周期管理onnxruntime 的 InferenceSession 不是线程安全的。如果服务端使用多线程并发处理翻译请求建议每个线程持有独立 session或者线程之间用锁保护session.run。我一开始图省事让所有线程共用一个 session结果线上偶发崩溃查了半天才定位到并发问题。更稳妥的做法是直接构建一个 session 池按请求量控制并发度。结尾一些真实经验真正做完这个项目之后我最大的感受是ONNX 迁移并不神秘难点不在“转换”本身而在于你要先理解模型的推理结构尤其是 Encoder-Decoder 和 KV Cache 之间的关系。只要掌握了这个基本功HuggingFace 上任何同类翻译模型都可以用同一套思路迁移到 ONNX。如果让我重新再做一遍我会先花 20 分钟把模型的 forward 流程完整过一遍再动手写导出脚本。第一次跑通的版本不用追求完美先 Greedy Search 出结果再上 Quantization 和 Beam Search逐步压榨性能。这个项目的收益非常直接模型小、速度快、部署干净、可控性强。我有一种很明确的预感以后我所有的 HuggingFace 小模型部署项目都会默认走一遍“转 ONNX 动态量化”这个流程。
返回列表