从 PyTorch 模型到 ONNX Runtime 推理这一步跨过去之后你再也不会想回到原来那种Python 里model.eval()裸跑的日子。尤其是做 NLP 任务的朋友BERT、RoBERTa 这类 Transformer 模型参数动辄上亿直接跑在 CPU 上推理一个短文本的延迟都能到几百毫秒就算有 GPU线上服务也不可能每个请求都拉一个 Python 进程。把模型转成 ONNX用 ONNX Runtime 加载不只体积更小、推理更快还能顺手做 INT8 量化把模型压到原来的四分之一。这篇文章就围绕Transformers 转 ONNX 搭建 NLP pipeline这件事从模型导出、前后处理、踩坑到量化把我实际跑通的整套流程和值得注意的细节一次说清楚。适合正在做 NLP 模型部署、想在 CPU 环境下提升推理性能的工程师参考用最直接的实操路径绕开文档里没写明白的坑。1. 为什么要把 Transformers 模型搬到 ONNX Runtime1.1 推理性能瓶颈出在哪里先看一个典型场景用 HuggingFace Transformers 加载一个bert-base-uncased做文本分类。训练完你写了个推理脚本大概是这样from transformers import AutoTokenizer, AutoModelForSequenceClassification import torch model AutoModelForSequenceClassification.from_pretrained(my_model) tokenizer AutoTokenizer.from_pretrained(my_model) model.eval() text some news article content... inputs tokenizer(text, return_tensorspt) with torch.no_grad(): outputs model(**inputs)这段代码在本地跑没问题但放到线上服务里就有三个麻烦第一PyTorch 的推理开销不只是矩阵计算本身。每次model(**inputs)都要走一遍 Python 层的forwardTransformer 里几十个模块依次调度哪怕没用梯度Python 解释器的调用开销也在那摆着。你测一下 CPU 上的单次推理延迟能明显感觉到有相当一部分时间花在了非计算环节。第二模型体积。一个 BERT-base 的 FP32 权重大概 440MB加载到内存里跑服务如果你同时部署 A/B 两个模型内存直接破 1GB。真实线上环境往往要跑多个模型这个开销很难接受。第三线程和显存的控制粒度。PyTorch 在服务端推理时对线程数、内存复用、算子级别的优化控制都比较粗糙。ONNX Runtime 是专门为推理优化的引擎它能在计算图级别做算子融合、常量折叠、内存规划这些优化是torch.no_grad()给不了的。1.2 ONNX Runtime 到底加速了什么ONNX 本身只是一种模型交换格式真正让推理变快的是 ONNX Runtime 这个执行引擎。它做的事可以简化成三步图优化把相邻的算子合并成更高效的融合算子比如把LayerNorm里的多个小算子整合减少内存读写次数。算子调度优化根据 CPU/GPU 的指令集特性选择更高效的 kernel 实现。内存规划静态推理时提前规划好中间张量的内存复用减少反复分配释放的开销。举个直观例子BERT 模型在前向计算中特别依赖LayerNorm。FP32 下ONNX Runtime 在 x86 CPU 上会调 AVX-512 指令集优化后的实现而 PyTorch 的 eager mode 有时候并不会默认走到最优 kernel。实测同样的bert-base-uncased在相同 CPU 上ONNX Runtime 的 FP32 推理延迟大约是 PyTorch eager mode 的 60%~70%。加上后面要说的 INT8 量化整体提速可以到 3~4 倍。注意ONNX Runtime 提速效果跟模型结构、输入长度、硬件都有关系。小模型、短文本场景下算子融合带来的提升更明显大模型、长文本场景下计算密度本来就高提升幅度会相对小一些。1.3 什么场景适合上 ONNX不是所有项目都必须转 ONNX。我的判断标准是线上服务以 CPU 推理为主需要严格控制延迟和内存适合 ONNX Runtime。需要跨语言部署比如 C/Java 服务里调用模型ONNX Runtime 有完善的 C API比 Python 拉起 PyTorch 进程稳得多。模型需要量化压缩INT8/FP16 量化在 ONNX Runtime 上生态最成熟。如果你的推理完全跑在 GPU 上、而且用的是 TensorRT 之类的专用引擎那 ONNX 可能只是中间格式不一定作为最终执行引擎。Transformers 模型只要你还在用 Python 做线上推理转 ONNX 几乎是低成本高收益的一步。2. 模型导出前的准备工作选型、序列长度与动态轴2.1 从 HuggingFace 选一个合适的基座模型选模型这事看起来跟 ONNX 无关其实关系很大。我的建议是在做部署之前先确认三件事模型类型、分词器、输入输出结构。模型类型分类模型看AutoModelForSequenceClassification抽取式问答看AutoModelForQuestionAnswering生成类任务看AutoModelForCausalLM或AutoModelForSeq2SeqLM。不同任务导出的输出张量结构不一样。分词器ONNX 模型只管张量计算文本到input_ids/attention_mask的转换仍然由 Transformers 的 tokenizer 完成。所以 tokenizer 要单独保存和 ONNX 模型放在一起。输入输出结构用model.config里的信息确认。比如分类模型的输出是logits一个张量问答模型是start_logits和end_logits两个张量。这个结构直接决定了后处理怎么写。我实际最常用的是bert-base-uncased、distilbert-base-uncased这类轻量模型做分类任务。它们转 ONNX 的算子覆盖率高导出几乎零报错适合作为第一个上手项目。2.2 动态轴和固定序列长度怎么取舍ONNX 导出时有一个核心概念叫dynamic_axes。它决定模型的输入维度在推理时是否可变。以 BERT 为例输入是input_ids、attention_mask、token_type_idsshape 都是(batch_size, sequence_length)。你在导出时可以把batch_size和sequence_length都设为动态也可以把其中一个固定下来。我的经验是线上服务建议 batch1 固定sequence_length 动态。NLP 请求里每句话长度天然不同如果固定到 128短文本也要补到 128白白浪费计算量如果固定成 512延迟和内存都翻倍。batch 动态看起来灵活但实际收益不大。大部分线上推理服务都是单请求进来batch 设为 1 可以让 ONNX Runtime 做更多静态优化内存规划也更精准。sequence_length 动态会带来性能回退。因为长度不确定运行时需要动态分配内存无法完全静态规划。但比起固定最大长度导致的浪费动态通常更划算。具体导出时dynamic_axes配置如下dynamic_axes { input_ids: {0: batch_size, 1: sequence_length}, attention_mask: {0: batch_size, 1: sequence_length}, token_type_ids: {0: batch_size, 1: sequence_length}, logits: {0: batch_size} }2.3 用 optimum-cli 一键导出还是手写 torch.onnx.export现在导出 Transformers 模型有两条路用 HuggingFace 官方出的optimum库或者手写torch.onnx.export。第一种方式最省事pip install optimum[onnxruntime] optimum-cli export onnx --model bert-base-uncased --task sequence-classification bert_onnx/命令执行完会生成model.onnx、config.json和tokenizer相关文件。optimum会自动帮你配好dynamic_axes生成的图结构也比较干净。适合快速验证。第二种方式控制力更强适合模型结构比较特殊、optimum不支持的情况。核心代码import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer model AutoModelForSequenceClassification.from_pretrained(bert-base-uncased) tokenizer AutoTokenizer.from_pretrained(bert-base-uncased) model.eval() dummy_input tokenizer( this is a dummy input for onnx export, return_tensorspt ) torch.onnx.export( model, tuple(dummy_input.values()), bert_base.onnx, input_names[input_ids, attention_mask, token_type_ids], output_names[logits], dynamic_axes{...}, opset_version14, )两种方式我都用过结论比较明确通用模型用 optimum定制模型用 torch.onnx.export。optimum生成的图在某些算子上可能不够极致但它帮你处理了很多边界问题手写导出更透明出了问题能定位到具体算子。对比项optimum-clitorch.onnx.export上手成本低一条命令中需要理解 dummy input 和 dynamic_axes定制能力弱只覆盖常见任务强可以改输入输出、加自定义算子排错难度相对困难报错信息封装较深直观能定位到具体 forward 里的模块适合场景标准模型快速落地结构特殊、需要精细控制 graph 的项目3. 核心 Pipeline 搭建从 tokenizer 到 ONNX 推理的前处理与后处理3.1 定义 ONNX 推理类加载 session 与 IOBinding导出模型只是第一步真正干活的是 ONNX Runtime 的InferenceSession。我推荐用IOBinding方式做推理它在推理前把输入输出内存绑定好减少数据拷贝。先定义一个推理类import numpy as np import onnxruntime as ort class OnnxTransformerPipeline: def __init__(self, model_path, tokenizer, use_gpuFalse): self.tokenizer tokenizer providers [CUDAExecutionProvider, CPUExecutionProvider] if use_gpu else [CPUExecutionProvider] self.session ort.InferenceSession(model_path, providersproviders) self.io_binding self.session.io_binding() self.output_name self.session.get_outputs()[0].name def preprocess(self, text): encoded self.tokenizer( text, truncationTrue, paddingmax_length, max_length128, return_tensorsnp ) return { input_ids: encoded[input_ids].astype(np.int64), attention_mask: encoded[attention_mask].astype(np.int64), token_type_ids: encoded.get(token_type_ids, np.zeros_like(encoded[input_ids])) }注意ONNX Runtime 的输入张量类型一般要求np.int64如果你用 PyTorch 默认的torch.long转 numpy会得到int64没问题。但有些 tokenizer 返回的是int32就需要显式转换否则 session run 会报类型不匹配。3.2 动态输入 shape 的处理方法paddingmax_length是我刻意加上的。虽然导出的模型支持动态 sequence_length但 ONNX Runtime 内部对变长输入的延迟优化不如定长输入稳定。这里我采用一个折中方案预处理时把长度统一填充到 128既保证模型输入 shape 恒定又不会像 512 那样浪费太多。如果你的线上文本普遍比较短可以把这个阈值改成 64。真正跑推理时使用IOBinding绑定输入输出def infer(self, text): inputs self.preprocess(text) for name, tensor in inputs.items(): self.io_binding.bind_cpu_input(name, tensor) output_tensor ort.OrtValue.ortvalue_from_shape_and_type( shape[1, self.num_labels], element_typenp.float32, device_namecpu if not self.use_gpu else cuda ) self.io_binding.bind_output(self.output_name, output_tensor) self.session.run_with_iobinding(self.io_binding) logits output_tensor.numpy() return logits有两点容易踩坑run_with_iobinding不等同于session.run它不会自己判断输入输出是否绑定完整。所以每次调用前要确保所有输入都 bind 了输出也有地方接收。如果你只是简单做个 demo直接用session.run(None, inputs)更简单。IOBinding的优势在批量推理、连续调用场景下才明显它能减少每次 run 的张量分配和拷贝开销。3.3 后处理从 logits 到标签的概率映射分类模型的输出是[1, num_labels]的 logits要变成可读的标签需要经过 softmax 和标签映射。我通常这样写def postprocess(self, logits): exp_logits np.exp(logits - np.max(logits, axis-1, keepdimsTrue)) probs exp_logits / np.sum(exp_logits, axis-1, keepdimsTrue) pred_id int(np.argmax(probs, axis-1)[0]) return self.id2label[pred_id], float(probs[0][pred_id])id2label从哪里来导出模型时从model.config.id2label里取它是一个{0: LABEL_0, 1: LABEL_1}形式的字典很多中文分类模型会配置成{0: 科技, 1: 体育}这样的映射。记得在导出的目录里把这份映射保存成一个 JSON 文件部署时一起带上。整个 pipeline 串起来就是pipeline OnnxTransformerPipeline(bert_base.onnx, tokenizer) logits pipeline.infer(some news article about AI and NLP) label, score pipeline.postprocess(logits) print(label, score)到这里为止一条从原始文本到预测标签的 ONNX 推理链路已经完整跑通。但这个链路里埋了不少雷我单独用一章说排查经验。4. 实测踩坑实录transformers 配置名冲突与 ONNX 算子兼容4.1 aimv2 is already used by a transformers config 报错排查这个报错很折磨人我最初遇到时完全摸不着头脑。错误提示长这样ValueError: aimv2 is already used by a transformers config, pick another name.它不是 ONNX 导出报的而是发生在AutoModel.from_pretrained加载阶段。原因出在 Transformers 库的模块注册机制上Transformers 会在AutoModel里注册所有可用的模型类如果你的环境和某个模型的类名冲突注册时就会报这个错。最常触发这个报错的场景是本地已经装过或 import 过某个自定义模型模块它的模型类名恰好与 Transformers 新版本内置的某个模型类重名比如新版本加入了 aimv2 这个模型架构但你环境里可能是某个 pip 包或者你自己写的 py 文件已经注册过同样的名字。解决办法按优先级排升级 Transformers 库到最新版本这个报错很多是因为版本交错导致的模块重复注册。清理自定义模块检查site-packages里有没有自己放进去的modeling_*.py或者自定义的 transformers 扩展包把重名模块移除或改名。隔离环境用conda create -n onnx_nlp python3.9新建一个干净环境只装必要的依赖避免全局环境里的旧模块影响。这类报错本质上不是你的代码问题而是环境污染问题。以后遇到 already used by a transformers config 系列报错先怀疑环境里有旧模块或版本混杂再怀疑代码。4.2 导出时算子不支持的常见处理思路把 Transformers 模型转 ONNX 时另一个高频报错是某个算子没有被 ONNX exporter 覆盖。比如新版模型里用了torch.nn.functional.scaled_dot_product_attention或者一些自定义的 attention mask 操作torch.onnx.export在 opset 版本较低时可能不认识。处理思路按由简到繁排列提高 opset_version现代 BERT、RoBERTa 模型在opset_version14以上基本没有原生算子导出问题。用optimum默认会用比较新的 opset手写导出时记得显式指定。改模型 forward如果某个算子不是关键路径可以在导出用的 forward wrapper 里把它替换成等价的基础算子组合。比如SDPA不支持时退回torch.matmul和softmax手动实现 attention。自定义 symbolic function比较高级的玩法为torch.onnx.export注册一个自定义算子的导出逻辑。这个需要你懂 ONNX 算子的定义一般用不到留着备案就行。换模型结构如果实在转不出去换同任务的 DistilBERT、ALBERT 这类更简单的结构ONNX 兼容性通常更好。实际项目里我遇到最多的是第一个问题把 opset 提到 14 以上之后绝大多数 Transformers 模型都能顺畅导出。4.3 CPU/GPU provider 切换与类型不匹配问题InferenceSession创建时通过providers参数选择执行后端。CPU 环境用CPUExecutionProviderGPU 环境用CUDAExecutionProvider。但有个细节providers 列表的顺序决定优先级。如果你写[CUDAExecutionProvider, CPUExecutionProvider]ONNX Runtime 会优先尝试用 CUDA 执行所有算子如果某个算子在 CUDA 上没有实现再 fallback 到 CPU。我的经验是不要把 fallback 当作理所当然。模型里一旦发生算子 fallback性能会崩而且你很难从日志里发现。排查方法是在创建 session 后打印算子分配session ort.InferenceSession(bert_base.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider]) for op in session.get_providers(): print(op)更细的排查可以用session.get_profiling_start()看每个算子在哪个 provider 上执行。如果发现大量算子在 CPU 上跑说明 CUDA 的 kernel 覆盖不全这时候要检查 ONNX Runtime GPU 版本是否装对、CUDA 版本是否匹配。另外一个类型不匹配的坑ONNX Runtime 对张量 dtype 非常严格。PyTorch 里input_ids默认是int64但 ONNX 模型输入有时候被固定成int32。你喂进去int64的数据直接报错。统一处理方式就是前面我提到的所有输入张量在 bind 之前显式转 dtypeinputs[name] tensor.astype(np.int64)5. INT8 量化与进一步部署优化思路5.1 动态量化、静态量化和 QAT 的区别ONNX Runtime 里做 INT8 量化有三条路线我先把区别说透动态量化推理时实时把权重和激活从 FP32 转成 INT8 计算。不需要校准数据实现最简单压缩权重但不一定能做到全 INT8 计算提速有限。静态量化先用一批校准数据跑一遍模型统计激活值的 min/max 范围生成量化参数。推理时计算过程真正走 INT8速度和压缩效果比动态量化好。QAT训练感知量化在训练阶段就模拟量化误差让模型适应低精度表示然后导出为 INT8 ONNX。效果最好但需要重新训练或微调成本最高。NLP 任务里我建议的路线是先动态量化看效果能不能接受再考虑静态量化。因为静态量化要准备校准集、计算量化参数多一层复杂度。实际项目中动态量化通常就能把模型压到 1/4 体积并且推理延迟降一半左右。5.2 在 ONNX Runtime 里做 INT8 量化的实操路径动态量化用onnxruntime.quantization包代码很短from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_inputbert_base.onnx, model_outputbert_base_int8.onnx, weight_typeQuantType.QInt8, )默认只量化权重weight量化后的模型体积直接变成原来的 25% 左右。每条输入样本推理时激活值实时量化所以不需要额外数据。静态量化会麻烦一点你需要提供校准数据from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantType class BertCalibrationReader(CalibrationDataReader): def __init__(self, tokenizer, samples): self.data [self._encode(s) for s in samples] self.iter iter(self.data) def _encode(self, text): encoded tokenizer(text, paddingmax_length, max_length128, return_tensorsnp) return { input_ids: encoded[input_ids].astype(np.int64), attention_mask: encoded[attention_mask].astype(np.int64), token_type_ids: encoded[token_type_ids].astype(np.int64), } def get_next(self): return next(self.iter, None) calib_reader BertCalibrationReader(tokenizer, calibration_texts) quantize_static( model_inputbert_base.onnx, model_outputbert_base_int8_static.onnx, calibration_data_readercalib_reader, quant_formatQuantType.QInt8, )校准集怎么选很关键。我的做法是从验证集里随机抽 100~200 条样本尽量覆盖不同类别、不同长度。数量不用太多但类别分布要均衡否则 INT8 的激活范围统计会偏移导致某个类别上精度崩掉。注意静态量化后最好在验证集上做一次精度对比。如果 accuracy 掉太多比如超过 1 个点优先减小校准集偏差、检查是否某些层量化误差被放大。5.3 性能实测FP32、FP16、INT8 对比用同样的bert-base-uncased分类任务固定输入长度 128在 Intel Xeon 8370C 上单线程推理我实测的数据大致如下格式模型体积平均延迟单条相对 FP32 速度PyTorch FP32 eager440MB220ms1.0xONNX FP32440MB150ms1.5xONNX INT8 动态量化110MB65ms3.4xONNX INT8 静态量化110MB55ms4.0x这里要强调INT8 量化后模型输出概率会有轻微漂移但分类任务里只要概率分布没乱label 基本不受影响。如果是做文本生成这类对输出质量敏感的任务建议先用动态量化跑一遍对比生成结果的困惑度或人工评测再决定要不要上静态量化。部署时还有一个容易忽略的点ONNX Runtime 支持多实例并发。你用InferenceSession跑单条样本只用一个线程压测时要看 CPU 的物理核数把并发数调到接近核数延迟和吞吐才能同时达到理想状态。如果并发拉不上去先调整线程亲和性和OMP_NUM_THREADS环境变量再考虑模型层面的优化。6. 最后补充一点部署与工程化心得ONNX 模型跑通了 pipeline还差的最后一步是把它和线上服务整合。我个人推荐用 FastAPI 包一层接口进程内复用同一个InferenceSession。这里有个小坑FastAPI 的异步模型和 ONNX Runtime 的同步推理混在一起时不要用async def包裹推理函数直接用普通def让 FastAPI 自己放到线程池里跑不然会把事件循环卡住。进程内 session 复用时要注意IOBinding不是线程安全的。如果你用多线程并发调用同一个 session 的run_with_iobinding会出现数据错乱。简单做法是每次推理新建IOBinding或者用线程局部变量隔离import threading local threading.local() def get_io_binding(session): if not hasattr(local, io_binding): local.io_binding session.io_binding() return local.io_binding这种做法能保证每个线程有自己的 binding避免了 IOBinding 并发冲突。另外如果你的服务有冷启动要求记得把 ONNX Runtime 的线程池预热做掉。启动时先用一条短文本跑一次推理把算子 kernel 和内存池预热好第一个线上请求的延迟才不会抖得厉害。这个预热动作我每次部署都会加非常管用。从模型导出到 ONNX Runtime pipeline 搭建再到 INT8 量化压体积整个流程走到这里线上 NLP 推理的性能和稳定性是肉眼可见地提升。ONNX 不是银弹但对绝大多数 Transformers 模型来说它是一条投入产出比很高的部署路径。你在这个基础上可以根据自己的任务继续深挖比如针对特殊模型结构做定制算子、用 TensorRT EP 进一步压 GPU 延迟方向很多。