Transformers模型转ONNX Runtime：NLP推理部署与INT8量化实战

从 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 Runtime

1.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_tensors="pt") with torch.no_grad(): outputs = model(**inputs)

这段代码在本地跑没问题，但放到线上服务里就有三个麻烦：

第一，PyTorch 的推理开销不只是矩阵计算本身。每次model(**inputs)都要走一遍 Python 层的forward，Transformer 里几十个模块依次调度，哪怕没用梯度，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_ids，shape 都是(batch_size, sequence_length)。你在导出时可以把batch_size和sequence_length都设为动态，也可以把其中一个固定下来。

我的经验是：

• 线上服务建议 batch=1 固定，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_tensors="pt" ) 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_version=14, )

两种方式我都用过，结论比较明确：通用模型用 optimum，定制模型用 torch.onnx.export。optimum生成的图在某些算子上可能不够极致，但它帮你处理了很多边界问题；手写导出更透明，出了问题能定位到具体算子。

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_gpu=False): self.tokenizer = tokenizer providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if use_gpu else ["CPUExecutionProvider"] self.session = ort.InferenceSession(model_path, providers=providers) self.io_binding = self.session.io_binding() self.output_name = self.session.get_outputs()[0].name def preprocess(self, text): encoded = self.tokenizer( text, truncation=True, padding="max_length", max_length=128, return_tensors="np" ) 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 的处理方法

padding="max_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_type=np.float32, device_name="cpu" 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, keepdims=True)) probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True) 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 文件）已经注册过同样的名字。

解决办法按优先级排：

1. 升级 Transformers 库到最新版本，这个报错很多是因为版本交错导致的模块重复注册。
2. 清理自定义模块：检查site-packages里有没有自己放进去的modeling_*.py，或者自定义的 transformers 扩展包，把重名模块移除或改名。
3. 隔离环境：用conda create -n onnx_nlp python=3.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_version=14以上基本没有原生算子导出问题。用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 环境用CPUExecutionProvider，GPU 环境用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 之前显式转 dtype：

inputs[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_input="bert_base.onnx", model_output="bert_base_int8.onnx", weight_type=QuantType.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, padding="max_length", max_length=128, return_tensors="np") 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_input="bert_base.onnx", model_output="bert_base_int8_static.onnx", calibration_data_reader=calib_reader, quant_format=QuantType.QInt8, )

校准集怎么选很关键。我的做法是从验证集里随机抽 100~200 条样本，尽量覆盖不同类别、不同长度。数量不用太多，但类别分布要均衡，否则 INT8 的激活范围统计会偏移，导致某个类别上精度崩掉。

注意：静态量化后，最好在验证集上做一次精度对比。如果 accuracy 掉太多（比如超过 1 个点），优先减小校准集偏差、检查是否某些层量化误差被放大。

5.3 性能实测：FP32、FP16、INT8 对比

用同样的bert-base-uncased分类任务，固定输入长度 128，在 Intel Xeon 8370C 上单线程推理，我实测的数据大致如下：

这里要强调，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 延迟，方向很多。