用过Transformers做NLP服务的人应该都经历过这种场景模型选的是BERT或者RoBERTa这类经典预训练模型推理时在CPU上跑延迟总是压不下来。上线一个文本分类接口单条请求动辄几十毫秒并发一高CPU直接被打满。GPU当然能解决一部分问题但成本摆在那里很多内部工具、边缘节点甚至客户现场的服务器根本配不起GPU。于是大家开始找各种加速方案ONNX Runtime就是这么进入了视线。这篇文章要聊的就是用ONNX搭建NLP Transformers pipelines的完整方案。ONNX是一个开放的模型交换格式它本身不是推理框架而是定义了一套统一的模型表示配合微软的ONNX Runtime推理引擎可以在CPU上拿到相当可观的性能提升。和原生的PyTorch推理相比ONNX Runtime在算子融合、内存复用、指令集优化上做了大量工作实测常见Transformer模型的推理延迟能缩短30%-60%再叠加int8量化很多场景能跑到接近GPU的水平。这篇文章面向的读者是那些已经会用Transformers库加载模型、跑通pipeline但还没真正把模型导出成ONNX做过优化的人。内容会覆盖从模型导出、pipeline搭建、int8量化到线上部署的完整链路最后还有我在实际项目中踩过的坑和排查经验。看完你就能自己动手把一条Transformers pipeline从PyTorch原生推理迁移到ONNX Runtime上并且知道每一步为什么要这么做。1. 先想清楚为什么Transformers pipelines需要ONNX1.1 原生pipeline的性能瓶颈在哪HuggingFace的Transformers库很香几行代码就能加载一个预训练模型处理文本分类、命名实体识别、问答等任务。但它的高性能推理能力其实并不突出。model AutoModelForSequenceClassification.from_pretrained(...)这种写法加载的是PyTorch的nn.Module对象推理流程完全走PyTorch的eager execution模式。eager execution的设计目标是灵活每一层都是动态执行Python解释器逐层调度。代价就是算子之间缺乏融合比如BERT里最常见的LayerNorm在PyTorch执行时会被拆成多个细颗粒度的小算子每个算子都要经过调度、显式分配内存、写回中间结果再传给下一个算子。这种碎片化的执行模式在GPU上有CUDA内核并行可以掩盖一部分但在CPU上线程同步和内存访问的开销会被放大得很明显。还有一点PyTorch的CPU推理默认走的是一套通用实现它要考虑跨平台兼容、动态shape支持、autograd的钩子等等。这些特性在训练阶段很关键但纯推理场景里全都是多余开销。ONNX Runtime的目标就是推理它不需要反向传播不需要动态图只需要把一个静态的计算图跑得足够快所以它可以把很多训练用的包袱丢掉专心做算子融合和内存规划。1.2 ONNX Runtime的加速原理拆解ONNX Runtime的加速能力核心可以归成四块。第一是算子融合。它会把计算图里的相邻算子合并成一个大算子。比如Transformer里常见的MatMul Add Layernorm Gelu这种组合在ONNX Runtime里会被融合成一个FusedGelu或Layernorm算子。算子数量少了内存读写次数就能降下来缓存命中率自然就上去了。第二是图优化。ONNX Runtime在加载模型时会做一轮常量折叠、冗余节点消除、维度推导。模型里有大量的权重矩阵这些是常量很多reshape、transpose操作可以直接在加载阶段就完成不用跑到推理时才做。第三是线程调度和内存池化。PyTorch推理时每个算子都可能会调用内存分配器ONNX Runtime则会在Session初始化时一次性规划好中间张量的大小复用内存块。配合OpenMP线程池在多核CPU上对大batch的并行推理效率提升非常明显。第四是底层指令集优化。ONNX Runtime的CPU实现大量使用了针对AVX、AVX2、AVX512指令集的手写优化矩阵乘法这块直接用了一流的BLAS实现比自己调用PyTorch的算子要高效得多。1.3 整体方案技术选型在这个方案里技术栈大概是这样的模型来自HuggingFace Transformers负责加载预训练权重和tokenizer用torch.onnx.export把PyTorch模型导出为ONNX格式推理阶段用onnxruntime的Python接口加载ONNX模型。整套链路不涉及C编译纯Python就能完成。有人可能会问Transformers官方不是也有optimum库支持导出ONNX吗这个库确实能用但它封装得比较重配置项不够透明一旦遇到不支持的算子排查起来比手动导出麻烦得多。我自己的习惯是先手动导出一次掌握整个流程的细节等模型结构简单、配置清晰时再用optimum做批量化导出也不迟。文章后面的示例都是用手动导出方式步骤可控性更强也方便大家理解底层逻辑。2. 环境准备与模型导出2.1 版本搭配和安装第一步是装依赖。需要注意ONNX Runtime、ONNX、Transformers、PyTorch这几个库的版本有兼容关系装得太乱容易出现莫名其妙的算子不支持。我自己在用的版本组合是pip install torch2.1.0 pip install transformers4.36.0 pip install onnx1.15.0 pip install onnxruntime1.17.1如果机器上有NVIDIA GPUONNX Runtime建议装onnxruntime-gpuCPU机器装onnxruntime就可以。这里有个容易踩的坑onnxruntime和onnxruntime-gpu不能同时装否则会互相覆盖Session初始化时会报冲突。装完之后可以用下面这行代码确认安装成功python -c import onnxruntime; print(onnxruntime.__version__, onnxruntime.get_available_providers())CPU设备应该能看到CPUExecutionProviderGPU设备能看到CUDAExecutionProvider。2.2 导出前需要理清的事情在写导出代码之前有三个问题需要先想清楚模型输入有哪些、输出是什么、哪些维度是动态的。以BERT分类模型为例输入通常有三个input_idstoken索引、attention_mask掩码、token_type_ids片段标识。在Transformers库中这三个张量虽然都有默认值但导出ONNX时必须显式地把它们作为模型输入传进去否则导出的模型在推理时拿不到完整的输入。然后是动态维度。绝大多数生产环境里一个batch内每条样本的长度是不同的模型需要支持变长序列输入。所以导出时必须把seq_len维度设为动态。同理如果服务接口支持batch推理batch_size维度也要设为动态。最后是固定参数。像max_length这类在tokenizer阶段就确定的值在模型内部不应该再出现。模型的输入张量shape只关心(batch_size, seq_len)具体长度是tokenizer处理时决定并向模型输入的。2.3 完整导出流程下面这段代码以bert-base-uncased作为演示模型导出一个文本分类任务的ONNX文件。核心是构造一个假输入然后调用torch.onnx.export。import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer model_name bert-base-uncased model AutoModelForSequenceClassification.from_pretrained(model_name, num_labels2) model.eval() # 构造一个形状吻合的假输入batch_size1, seq_len128 dummy_input_ids torch.randint(0, 1000, (1, 128), dtypetorch.long) dummy_attention_mask torch.ones(1, 128, dtypetorch.long) dummy_token_type_ids torch.zeros(1, 128, dtypetorch.long) torch.onnx.export( model, (dummy_input_ids, dummy_attention_mask, dummy_token_type_ids), bert-base-uncased.onnx, input_names[input_ids, attention_mask, token_type_ids], output_names[logits], dynamic_axes{ input_ids: {0: batch_size, 1: seq_len}, attention_mask: {0: batch_size, 1: seq_len}, token_type_ids: {0: batch_size, 1: seq_len}, logits: {0: batch_size}, }, opset_version17, do_constant_foldingTrue, )几个参数值得多说两句。opset_version代表ONNX算子集的版本ONNX Runtime 1.17对opset 17支持得非常完整低于13有些新算子不支持太高则可能遇到编译器兼容问题。do_constant_foldingTrue会让PyTorch在做torch.jit.trace时就把常量折叠掉很多权重变换在导出阶段就完成了减少推理期的计算量。导出成功后目录下会多出一个bert-base-uncased.onnx文件你可以用onnx.checker.check_model验证结构的合法性import onnx model onnx.load(bert-base-uncased.onnx) onnx.checker.check_model(model) print(模型结构合法)2.4 导出后必须做的两步验证ONNX模型合法不代表它拿到的推理结果和PyTorch一致。这一步至关重要很多人导出来之后直接上线结果预测分数和原来不一样还找不到原因。先用ONNX Runtime加载模型和PyTorch跑同一批输入做对比import numpy as np import onnxruntime as ort session ort.InferenceSession(bert-base-uncased.onnx, providers[CPUExecutionProvider]) # 同一份输入 text This film is really good! inputs tokenizer(text, return_tensorspt, paddingTrue, truncationTrue, max_length128) # PyTorch结果 with torch.no_grad(): pt_logits model(**inputs).logits.numpy() # ONNX结果注意要转成 numpy数组 并指定 dtype onnx_inputs { input_ids: inputs[input_ids].numpy().astype(np.int64), attention_mask: inputs[attention_mask].numpy().astype(np.int64), token_type_ids: inputs[token_type_ids].numpy().astype(np.int64), } onnx_logits session.run([logits], onnx_inputs)[0] # 对比 diff np.abs(pt_logits - onnx_logits).max() print(f最大绝对误差: {diff:.6f})如果误差在1e-4量级以内基本可以认为导出没问题。如果误差到了1e-1甚至更大大概率是模型没有切到eval模式、dropout还在生效或者动态轴的配置有误导致某些shape信息不对。另外提醒一句PyTorch端的torch.no_grad()一定要加。如果不加模型里的dropout层会随机丢弃节点导出的ONNX模型和推理时的PyTorch模型行为不一致对比时误差会忽大忽小。3. 搭建ONNX Runtime推理Pipeline3.1 体验一次最朴素的ONNX推理模型导好了接下来做推理。最直接的方法是手动管理输入输出代码很短但对理解机制很有帮助import numpy as np import onnxruntime as ort session ort.InferenceSession(bert-base-uncased.onnx, providers[CPUExecutionProvider]) def predict(text: str): inputs tokenizer( text, return_tensorsnp, # 返回numpy省去转换 paddingTrue, truncationTrue, max_length128, ) feeds { input_ids: inputs[input_ids].astype(np.int64), attention_mask: inputs[attention_mask].astype(np.int64), token_type_ids: inputs[token_type_ids].astype(np.int64), } logits session.run([logits], feeds)[0] return logits result predict(the movie is fantastic!) print(result)注意一个细节return_tensorsnp会让tokenizer直接返回numpy数组省去了从PyTorch张量转numpy的时间。这个细节在高并发场景下也有意义tokenizer阶段省下的每一点开销都会累计到整体延迟里。ONNX Runtime的session.run第一个参数是输出节点名列表如果需要多个输出可以一次传多个名字第二个参数是输入字典key必须和导出时的input_names完全一致否则会报Invalid Feed Input Name的错误。3.2 升级成可复用的Pipeline类手动管理输入输出写两次还行要是服务里要接十几个不同的模型就会想封装了。我习惯写一个通用的OnnxPipeline类把tokenizer、session、任务类型都收进来对外只暴露一个__call__方法。class OnnxPipeline: def __init__(self, onnx_path: str, tokenizer, task: str text-classification): self.session ort.InferenceSession(onnx_path, providers[CPUExecutionProvider]) self.tokenizer tokenizer self.task task def _prepare_inputs(self, text: str): encoded self.tokenizer( text, return_tensorsnp, paddingTrue, truncationTrue, max_length128, ) 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 __call__(self, text: str): feeds self._prepare_inputs(text) logits self.session.run([logits], feeds)[0] if self.task text-classification: probs np.exp(logits - logits.max(axis-1, keepdimsTrue)) probs probs / probs.sum(axis-1, keepdimsTrue) return {label: int(probs.argmax(axis-1)), score: float(probs.max(axis-1))} return logits实际项目里还要考虑batch推理。对多个文本一起处理时tokenizer会在paddingTrue的情况下把较短样本补齐到batch内最长长度这样ONNX模型可以直接吃到(batch_size, max_seq_len)的输入一次推理完成多个样本。batch推理时延迟的增长是亚线性的4条样本一起跑通常比逐条跑4次要快不少。3.3 扩展到问答、NER等任务的差异文本分类的输入输出是最简单的情况但实际业务里更多是NER、抽取式问答这类任务它们的输出结构不一样导出时也要做相应调整。NER任务一般用AutoModelForTokenClassification输出的是每个token的标签分布shape是(batch_size, seq_len, num_labels)。导出时输出张量是logits但shape里包含seq_len所以dynamic_axes里要把logits的0和1两个维度都标成动态。问答任务用AutoModelForQuestionAnswering输出不仅仅是logits还有start_logits和end_logits两个张量。导出时可能需要一次性导出两个输出torch.onnx.export( model, (dummy_input_ids, dummy_attention_mask, dummy_token_type_ids), qa_model.onnx, input_names[input_ids, attention_mask, token_type_ids], output_names[start_logits, end_logits], dynamic_axes{ input_ids: {0: batch_size, 1: seq_len}, attention_mask: {0: batch_size, 1: seq_len}, token_type_ids: {0: batch_size, 1: seq_len}, start_logits: {0: batch_size, 1: seq_len}, end_logits: {0: batch_size, 1: seq_len}, }, opset_version17, )推理拿到两个输出后还需要做一层后处理对start_logits和end_logits做softmax然后在有效的span范围内找start和end得分最高的组合再映射回原始文本的字符偏移量。这一部分用Transformers库里的postprocess_qa_predictions可以直接搞定不用自己从零写。3.4 服务化部署时的输入输出约定ONNX模型一旦投入服务化环境就需要显式定义输入输出的契约。和直接用Transformers pipeline不同ONNX模型的输入是张量没有tokenizer的处理逻辑所以服务层的输入格式设计很重要。我常用的做法是HTTP API层统一接收原始文本然后在一个薄封装层里调用tokenizer把原始文本变换成张量再交给ONNX模型。这样模型服务本身只负责张量到张量的计算没有任何字符串处理的负担是纯计算型微服务。tokenizer和模型的关系被固定在这个封装层里不会出现线上用的tokenizer和导出模型时用的tokenizer不一致的问题。还要注意Transformers库的tokenizer版本更新后词汇表内容有时会变化导致推理结果和之前不一样。这类问题非常隐蔽因为它不会报错但准确率就是莫名掉了。所以模型文件和tokenizer文件必须绑定发布建议把它们放在同一个版本目录下并且在上线后跑一遍回归测试。4. 性能优化int8量化与运行时配置4.1 先搞清楚ONNX Runtime自己的优化选项很多人在模型导出后就直接开始量化其实ONNX Runtime本身就提供了相当大程度的图优化能力。InferenceSession可以设置sess_options其中graph_optimization_level是影响性能的一个重要参数。sess_options ort.SessionOptions() sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.intra_op_num_threads 4 sess_options.inter_op_num_threads 2 session ort.InferenceSession( bert-base-uncased.onnx, sess_optionssess_options, providers[CPUExecutionProvider], )ORT_ENABLE_ALL会开启所有图优化包括算子融合和布局优化。intra_op_num_threads控制单个算子内部的线程数inter_op_num_threads控制算子间并行度。对于BERT这类模型实测下来单算子内部并行收益更高所以通常intra_op给足线程inter_op保持较小值。如果在多路物理CPU的机器上部署还要注意设置sess_options.affinity_mask或使用OpenMP环境变量控制线程绑核。不控制绑核的话操作系统可能频繁切换线程性能波动很大。4.2 动态量化实操ONNX Runtime的量化分为动态量化和静态量化。动态量化不需要校准数据实现成本最低适合快速上手。它的原理是只把权重从float32量化到int8激活值在推理时动态统计范围并量化计算完再反量化回float32。好处是精度损失通常较小坏处是激活值的量化/反量化过程本身有一定开销。用quantize_dynamic就够两三行代码from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( bert-base-uncased.onnx, bert-base-uncased-int8.onnx, weight_typeQuantType.QInt8, )weight_type可选QInt8和QInt4QInt8通用性最好QInt4可以把模型压得更小但目前对CPU的优化支持不如QInt8成熟。模型体积直接从120MB缩到40MB左右对容器镜像体积和冷启动速度都是很大改善。量化完一定要重新跑一遍精度对比。BERT类模型在做文本分类时动态量化带来的准确率下降通常不超过0.5%但如果是做回归任务或者对score特别敏感的业务就必须用验证集仔细评估。4.3 静态量化的校准流程静态量化比动态量化更激进它会把激活值也量化为int8推理时不再动态计算范围而是使用预设的量化范围。这样推理速度更快但对校准数据的依赖很重。流程分为三步准备数据、收集范围、量化。ONNX Runtime提供了quantize_static接口配合CalibrationDataReader使用。from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantType class DataReader(CalibrationDataReader): def __init__(self, samples): self.samples samples self.iter iter(samples) def get_next(self): return next(self.iter, None) calib_samples [] for text in calib_texts: inputs tokenizer(text, return_tensorsnp, paddingTrue, truncationTrue, max_length128) calib_samples.append({ input_ids: inputs[input_ids].astype(np.int64), attention_mask: inputs[attention_mask].astype(np.int64), token_type_ids: inputs[token_type_ids].astype(np.int64), }) quantize_static( bert-base-uncased.onnx, bert-base-uncased-static-int8.onnx, DataReader(calib_samples), weight_typeQuantType.QInt8, )校准集的选择直接影响量化质量。理想情况是贴近真实业务分布的样本多样性要够覆盖长文本、短文本、不同风格的内容。校准集样本量一般在几百到几千条之间太少会导致激活值范围估计不准太多则校准耗时长但收益不明显。静态量化在CPU上的推理延迟通常比动态量化再低20%-40%但精度风险也更高。实际操作中如果静态量化精度掉得厉害可以退回到动态量化或者对特定层跳过量化。4.4 实测数据对比我在一台8核Xeon CPU无AVX512上用bert-base-uncased做过一轮对比单条文本平均长度为64 tokenbatch1结果大致如下配置平均延迟(ms)模型体积(MB)相对基线PyTorch CPU推理28.6420含PyTorch1.0xONNX Runtime FP3212.41202.3xONNX Runtime 动态int86.8404.2xONNX Runtime 静态int85.1405.6x这个数据仅供大家参考不同硬件、不同序列长度、不同batch size下的收益会不一样。但从趋势上看ONNX Runtime在CPU推理上的优势是实打实的尤其int8量化后性能提升非常直观。还有一点值得提的是batch推理的收益。同样是ONNX Runtime动态int8模型batch8的延迟约15ms平均到每条样本只有1.9ms。所以如果业务允许合并请求用batch推理能进一步把吞吐拉起来。5. 常见问题排查与避坑实录5.1 导出阶段的高频报错导出阶段最常见的报错之一是Unsupported operator。某些模型里用了自定义的attention实现PyTorch在trace时没法映射到ONNX算子就会报这个错。解决办法是换一个标准的模型实现或者升级opset版本让更多算子被支持。另一个思路是用torch.onnx.export的operator_export_type参数切换导出模式以兼容更多算子组合。shape inference failed的报错也经常出现一般是因为dynamic_axes配置不完整导致ONNX的shape推断失败。解决方法是把模型中所有涉及seq_len维度的输入输出都加到dynamic_axes里确保每个张量的动态维度都被声明。还有一类问题不是报错而是导出成功但模型文件巨大甚至达到GB级别。这种通常是导出时把整个优化器和训练状态都带进去了检查一下导出模型的结构重点看权重名称里是否有optimizer、scheduler之类的节点如果有说明导出方式有问题应该用model AutoModelForXxx.from_pretrained(...)加载后直接导出而不是加载训练器再导出。5.2 推理阶段的精度与性能问题推理阶段最典型的诡异问题是ONNX模型的单条推理结果和PyTorch一致但batch推理时结果不同。原因通常出在tokenizer的padding策略上。Transformer的attention mask在padding部分应该为0但如果某个模型导出时没有正确处理attention mask比如模型实现里没把它当输入padding位置也会参与计算导致结果出现偏差。排查方法是逐条对比单条推理和batch推理时的logits通过观察差异分布来判断是不是padding部分导致的。性能不达标的情况也时有发生。如果ONNX Runtime推理延迟和PyTorch差不多首先要检查graph_optimization_level是否设成了ORT_ENABLE_ALL。然后确认线程数配置合理intra_op_num_threads不超过物理核数。最后检查是否用了量化的模型文件——有时量化文件生成没问题但加载时由于路径写错实际加载的还是FP32模型这种错误代码上不报错只有延迟值和模型体积能暴露问题。5.3 一个容易被忽略的config命名冲突有读者遇到过报错aimv2 is already used by a transformers config, pick another name.这个错误其实和ONNX没什么直接关系但常见于用Transformers保存或者导出模型时本地目录或者HuggingFace缓存目录里出现了重名的config条目。Transformers在加载某个预训练模型时会先查缓存和本地目录如果发现了两个相同的模型类型名称或者config名称就会抛出这个提示来提醒你清理冲突。解决方法是检查~/.cache/huggingface目录下是否有重名目录或者设置一个全新的cache_dir重新拉取。另外如果你用save_pretrained保存过自定义模型名建议显式传入一个唯一的config名称避免和官方仓库里的名字撞车。这个问题很底层但一旦踩到就是硬卡半天查不出原因。5.4 实战避坑清单最后整理一份我在多次实践中沉淀下来的避坑清单照着检查能省不少时间。导出前一定执行model.eval()并且用torch.no_grad()包住推理过程否则dropout会破坏导出的图结构。tokenizer的return_tensors最好统一为np既能喂给ONNX又避免和PyTorch张量混合时出现隐式转换。动态轴的名称必须和模型内部的输入张量名称完全一致大小写也对不上就会报feed错误。量化后的模型精度验证不要只跑一两条样本至少用一个100条以上的测试集统计整体指标变化。ONNX Runtime版本升级后务必重新跑一遍精度对比和性能测试算子融合策略可能会变延迟和精度都有波动。多线程下部署时注意线程池的配置和模型实例的数量避免出现线程饥饿或者CPU过度订阅。线上模型和tokenizer要版本联动升级任意一方都要做回归验证否则会出现“模型没变结果变差”的诡异现象。最后再分享一点个人体会ONNX这条路我走了很长时间最初的动机很简单就是想在CPU上把BERT推理延迟压进10毫秒以内。踩过不少坑之后回头看ONNX Runtime真正厉害的地方不只是推理引擎的算子优化更多是它逼着你去重新审视整个NLP服务的链路——从模型的动态维度设计、tokenizer和模型的契约、量化策略的选择到线程调优和部署形态每个环节都变得清晰可控了。也许你当前的项目只是给一个小工具加一个pipeline但ONNX这套方法论迁移到CV、TTS等其他模型上一样通用它的价值会随着模型场景的增多逐渐放大。希望这篇文章能让你少走一些弯路把加速做实而不是停留在“跑通了一个demo”的层面。