接到这个任务的时候团队里已经有一个在训练环境里跑得不错的VisionTransformer图像分类模型。模型是OpenMMLab工具链下训练出来的权重文件几百MB评估精度很高可真正把它部署到生产环境的CPU机器上用onnxruntime推理时问题就全冒出来了导出的ONNX里多了一堆无用算子动态维度怎么处理预处理一个细节没对齐精度立刻掉一个点。折腾完整个流程后我把这次的完整过程整理成一份记录希望能给正卡在同样环节的朋友一点参考。内容不复杂但每一步都有讲究环境版本、导出方式、预处理对齐、动态Shape、推理性能少看一个都可能让模型在线上跑出错误结果。1. 先把部署链路想明白OpenMMLab、ViT与ONNX Runtime各是什么角色1.1 VisionTransformer的部署难点在哪VisionTransformer和传统CNN最大的区别是它的前向过程里有大量张量变形操作。卷积网络里的卷积、池化、激活函数在ONNX里都有非常成熟的算子映射几乎不会出错而Transformer结构里的核心组件是Patch Embedding、Multi-Head Attention、LayerNorm、残差连接这些模块在PyTorch里写起来很顺手一旦要导出成ONNX就会出现很多特殊算子尤其是在不同OPSet版本下Reshape、Transpose、Unsqueeze、Gather这些操作的兼容性很容易出问题。举个实际的例子注意力机制里经常要执行q q * self.scale、attn softmax(q k.transpose(-2, -1) / sqrt(d))这里面的转置、除以常数、Softmax组合起来在ONNX里会生成一系列小算子。如果模型里还带了相对位置索引、可变形注意力这类自定义模块那导出时更加容易把某个自定义函数留在图里导致onnxruntime直接不支持。VisionTransformer模型的输入通常固定成224x224或者384x384这是因为位置编码Position Embedding的形状是训练时固定下来的。如果输入分辨率变了要么把Position Embedding做插值要么干脆用固定Shape导出。很多团队一上来就想做动态分辨率结果在推理时发现位置编码对不上最终精度崩掉这就是对模型结构理解不足导致的。OpenMMLab系列的图像分类仓库里已经收敛了ViT、DeiT、Swin等大量Transformer分类结构。它们统一用配置文件管理数据预处理、模型结构、损失函数和训练参数模型前向输出也被封装得相对规范。表面上看起来“OpenMMLab导出模型”只是一个命令行操作但真正的难点在于你要知道工具链帮你封装了哪些逻辑导出之后的模型接收什么输入输出什么语义以及怎么保证预处理和后处理和训练时完全一致。1.2 为什么最终选择ONNX Runtime作为执行后端部署模型有很多种方案PyTorch直接部署、TensorRT、OpenVINO、ONNX Runtime等。如果目标平台是Linux服务器且要求Python环境ONNX Runtime是最省事的选择。它提供了统一的推理API底层支持CPU和GPU不需要重新编译框架也不需要写C代码只要把模型转成ONNX几乎在任何支持ONNX Runtime的机器上都能跑。和TensorRT相比ONNX Runtime的优化虽然不如TensorRT那么激进但胜在简单稳定。TensorRT在上手ViT时通常会碰到固定Shape、算子插件、显存分配等一系列问题而ONNX Runtime只要模型图能被读取基本上跑起来不会出大问题。团队里如果希望快速验证部署效果onnxruntime是最合适的第一站。它还能通过ExecutionProvider方便地在CPU和GPU之间切换比如配置[CUDAExecutionProvider, CPUExecutionProvider]机器上有显卡就用GPU没有显卡就自动落到CPU这对生产环境非常有价值。选ONNX Runtime还有一个很重要的原因它能和OpenMMLab里的MMDeploy生态无缝衔接。MMDeploy提供了很多现成的ONNX Runtime配置模板我们只要在配置文件里指定后端类型转换脚本就会自动处理好预处理、输出格式和动态Shape的映射。后面我写步骤时会分别介绍两条路线一条是官方MMDeploy转换脚本另一条是直接用PyTorch的torch.onnx.export手工导出。两条路线我都跑通过各有适用场景。2. 环境搭建与模型初步验证2.1 版本配对是第一个隐形门槛OpenMMLab系列工具链版本断裂是出了名的。不同版本的mmcv、mmcls、mmengine之间有严格的对应关系如果直接pip install最新版常常会遇到某个模块找不到或者编译好的mmcv和本地PyTorch不兼容的问题。我这里踩过的坑先列出来如果使用mmcls0.25.0建议对应mmcv1.5.0和pytorch 1.8-1.10如果使用mmpretrain1.x建议对应mmcv2.0.0和pytorch 1.13以上。绝大多数兼容性问题都出在mmcv的编译上因为它包含C和CUDA算子安装时如果没有找到匹配的预编译包可能需要现场编译时间非常长。一个比较稳妥的安装流程是# 创建虚拟环境 conda create -n vit_deploy python3.8 -y conda activate vit_deploy # 安装PyTorch pip install torch1.13.1 torchvision0.14.1 --index-url https://download.pytorch.org/whl/cu117 # 安装mmcv和mmpretrain pip install -U openmim mim install mmcv2.0.0 pip install mmpretrain等这些基础依赖装完之后再用pip install onnx onnxruntime onnxsim安装导出和推理相关的库。onnxruntime这里有个小坑如果做CPU推理直接装官方onnxruntime即可如果要做GPU推理需要安装onnxruntime-gpu而且要确认它自带的CUDA版和本地CUDA版本能对应上否则初始化时非常容易直接崩溃。提示如果你在一个全新机器上部署建议不要一上来就用最新的mmpretrain。ai模型代码迭代非常快今天能跑通的配置文件过三个月可能就被标记为废弃了。锁定一个稳定版本比追求新功能更重要。2.2 权重和预训练模型拉下来先做一次前向对齐在开始导出之前一定要先在PyTorch环境里跑一遍模型拿到推理结果和中间张量这能帮你建立“基准答案”后续ONNX导出的输出要和它做精度对比。如果没有这一步导出后出错了都不知道是哪里引入的误差。以mmpretrain为例最简单的推理方式是调用inference_modelfrom mmpretrain.apis import inference_model, init_model config_path configs/deit/deit-tiny_4xb256_in1k.py checkpoint_path https://download.openmmlab.com/mmclassification/v0.3/deit/deit-tiny_4xb256_in1k_20220124-bf3d9f50.pth model init_model(config_path, checkpoint_path, devicecpu) result inference_model(model, demo.jpg) print(result[pred_class]) print(result[pred_scores])这个函数内部会自己处理数据预处理、模型前向、后处理所以你看到的是一个完整的分类结果。但如果后面要导出ONNX不能直接用这个封装好的模型对象导出因为inference_model返回的result里除了预测分数还有类别名、可视化数据这些不是模型前向的一部分直接导出会把这些后处理逻辑绕进去。比较规范的做法是拿到模型对象后用model(x, return_lossFalse)或者自定义一个只做前向的包装类来调用。这里需要特别注意mmcls / mmpretrain 的模型forward方法带有return_loss参数训练时传True会返回loss推理时传False才返回logits。我们在导出时一定要把包装逻辑写清楚否则导出的ONNX里可能混入loss计算分支。import torch from mmpretrain.apis import init_model model init_model(config_path, checkpoint_path, devicecpu) model.eval() x torch.randn(1, 3, 224, 224) with torch.no_grad(): logits model(x, return_lossFalse) print(logits.shape) # 应该是[1, num_classes]这一步跑通之后把logits保存成npy文件留到后面和ONNX Runtime的输出做对比。3. 导出ONNX两种可以落地的路线3.1 路线一用MMDeploy官方转换脚本MMDeploy是OpenMMLab体系里的部署工具它把模型转换过程做成了一套配置驱动的流程。针对图像分类任务MMDeploy提供了ONNX Runtime后端对应的转换配置使用方式很直接。首先把MMDeploy仓库拉下来安装它的依赖git clone https://github.com/open-mmlab/mmdeploy.git cd mmdeploy pip install -e .然后用它提供的tools/deploy.py转模型。命令大概是这样的python tools/deploy.py \ configs/mmcls/classification_onnxruntime_dynamic.py \ /path/to/deit-tiny_4xb256_in1k.py \ /path/to/deit-tiny_4xb256_in1k_20220124-bf3d9f50.pth \ /path/to/demo.jpg \ --work-dir work_dir/deit \ --device cpu这里配置文件名里如果带dynamic意味着导出的ONNX输入Shape是动态的如果是static则输入Shape固定。对ViT来说我建议先导出固定224x224的静态模型除非你确认自己需要动态尺寸且已经处理了位置编码。转换完成后在work_dir/deit下会生成deit.onnx文件和对应的部署配置文件。这个脚本的优点是省心它自动处理了return_loss分支、输出层解析以及数据预处理中的若干细节。缺点是调试空间比较少如果某个自定义算子不支持你只能去看它的源码再手动改配置。我的经验是先用这个脚本转一遍能成最好如果不能成再走下面的手工路线。3.2 路线二在PyTorch里包一层再torch.onnx.export如果你熟悉模型结构想完全控制导出参数那么手工导出路线反而更灵活。核心思路是在PyTorch里写一个包装模块只保留模型的前向推理路径然后调用torch.onnx.export。具体做法是import torch from mmpretrain.apis import init_model class ViTWrapper(torch.nn.Module): def __init__(self, model): super().__init__() self.model model def forward(self, x): return self.model(x, return_lossFalse) model init_model(config_path, checkpoint_path, devicecpu) model.eval() wrapper ViTWrapper(model).eval() onnx_path vit_tiny.onnx x torch.randn(1, 3, 224, 224) torch.onnx.export( wrapper, x, onnx_path, input_names[input], output_names[logits], opset_version13, do_constant_foldingTrue, dynamic_axes{ input: {0: batch_size}, logits: {0: batch_size} } )这里有两个参数需要注意。第一个是opset_version我通常用11或13。opset 11对旧版onnxruntime兼容性最好opset 13对某些新算子的支持更好但如果代码里用了比较新的算子可能导致旧版onnxruntime加载失败。第二个是dynamic_axes我在这里只把第0维设为动态也就是batch维度可变而高和宽保持固定。ViT模型的位置编码是固定维度的如果你把高宽也设为动态虽然可以导出来但实际推理时如果输入尺寸不是224x224位置编码的插值是个大麻烦。包装类的forward里传入return_lossFalse这个动作看着小作用却是决定性的它把模型从训练模式切换到推理模式跳过loss计算和辅助分支输出才是干净的logits。3.3 导出后的模型检查与简化不管是哪条路线导出的模型强烈建议都做一遍结构检查和简化。简化工具使用onnxsim它能合并一些常量节点、删除冗余计算让图的规模显著缩小。ViT这种模型天然包含大量reshape和transpose很多中间节点都是常量折叠可以消除的用不用onnxsim推理速度差距可以达到20%以上。python -m onnxsim vit_tiny.onnx vit_tiny_sim.onnx \ --overwrite-input-shape 1,3,224,224执行成功后可以用onnxruntime先做一个简单推理同时用onnx.shape_inference.infer_shapes检查每个节点的Shape信息是否完整。我习惯用Netron打开简化后的ONNX文件先看输入节点、输出节点再把注意力模块里的节点大致过一遍确认有没有红框报错的算子。一旦看到红色节点就说明这个算子在onnxruntime里可能不支持越早发现越好解决。注意ONNX虽然是统一的中间格式但ONNX Runtime对不同算子的支持程度并不完全等于onnx标准。同一个模型用Netron打开完全正常一到onnxruntime的CPU上就报错的情况很常见。所以在导出后立即用onnxruntime加载并推理一次比做任何静态检查都靠谱。4. 图像预处理和后处理精度能不能对齐的关键4.1 把mmcls预处理拆开来看很多人在部署时只关注模型转换对不对忽略预处理结果推理出来的分类结果和训练时完全对不上。ViT训练时的pipeline通常包含Resize到固定尺寸然后CenterCrop到网络输入尺寸再做Normalize。在mmpretrain配置里一般长这样train_pipeline [ dict(typeLoadImageFromFile), dict(typeResize, size256), dict(typeCenterCrop, crop_size224), dict(typePackInputs), ]这里的细节非常多。Resize到256再CenterCrop到224这个“先把短边放大到256再从中心裁224”的操作顺序不能改。如果直接Resize到224分类精度通常会下降一到两个百分点。Normalize时用的均值和标准差是ImageNet上的标准统计量mmpretrain默认是mean[0.485, 0.456, 0.406]、std[0.229, 0.224, 0.225]。这个值要体现在onnxruntime的代码里不能漏。还有一个最常踩的坑是通道顺序。OpenCV读图是BGRPIL读图是RGBPyTorch训练时一般用PIL读图然后转tensor所以通道顺序是RGB。如果用onnxruntime推理时拿OpenCV读图却不转成RGB模型输出的置信度会乱成一团。正确做法是读图后用cv2.cvtColor(img, cv2.COLOR_BGR2RGB)或者干脆用PIL读取。归一化的公式也有讲究PyTorch里ToTensor会把像素值从0到255缩放到0到1然后Normalize再执行(x - mean) / std。所以在onnxruntime的预处理里不能光减均值再除方差必须先除以255。有些部署代码图省事把均值和方差直接乘上255再作减除理论上等价但更容易写错不如老老实实分步写。4.2 后处理与top-5统计导出ONNX时模型输出的是logits不是概率。如果要得到置信度需要自己加一层Softmax。因为模型head里通常只输出全连接层的logits后处理加上Softmax这样和PyTorch里原来的predict结果才能对上。在onnxruntime推理时后处理代码可以这样写import numpy as np def postprocess(logits, topk5): scores logits - logits.max(axis-1, keepdimsTrue) exp_scores np.exp(scores) probs exp_scores / exp_scores.sum(axis-1, keepdimsTrue) topk_idx np.argsort(probs)[0][-topk:][::-1] topk_score probs[0][topk_idx] return topk_idx, topk_score手动实现Softmax是为了确保数值稳定因为ViT的logits范围可能比较大直接指数容易溢出。这样得到的top-5结果如果和PyTorch里inference_model的pred_score进行比较误差在千分之一以内都算正常。提示模型输出logits可能是[1, 1000]也可能因为训练时用了其它数据集而变成[1, N]。不要在生产代码里硬编码1000最好根据session.get_outputs()[0].shape动态获取类别数。这样以后替换模型时不用改代码。5. onnxruntime推理实现与性能实测5.1 搭建CPU/GPU可切换的推理会话onnxruntime的会话配置是影响性能和稳定性的重要环节。最基本的加载方式是一句ort.InferenceSession但要在生产环境使用建议使用下面的配置方式import onnxruntime as ort sess_options ort.SessionOptions() sess_options.intra_op_num_threads 8 sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.log_severity_level 3 providers [CUDAExecutionProvider, CPUExecutionProvider] session ort.InferenceSession( vit_tiny_sim.onnx, sess_optionssess_options, providersproviders )intra_op_num_threads控制每个算子内部的线程数。对ViT这种结构来说很多算子都是小矩阵乘法线程数开太高反而会因为线程切换增加开销。一般CPU机器开4到8线程比较合适具体可以在机器上扫描一遍找到最优值。graph_optimization_level开启全部优化其中包含一些算子融合对Transformer模型通常有正面效果。providers列表的顺序有讲究onnxruntime会按顺序尝试可用的Provider第一个不可用就尝试下一个。所以把GPU放前面能在有显卡的机器上自动用GPU没有显卡就回落CPU。5.2 完整推理代码下面给出一份可以在生产环境使用的完整推理代码输入一张图片路径输出top-5结果import numpy as np from PIL import Image import onnxruntime as ort MEAN np.array([0.485, 0.456, 0.406], dtypenp.float32).reshape(1, 3, 1, 1) STD np.array([0.229, 0.224, 0.225], dtypenp.float32).reshape(1, 3, 1, 1) def load_and_preprocess(image_path, input_size224): img Image.open(image_path).convert(RGB) # Resize to 256 then center crop to 224 img img.resize((256, 256), Image.BILINEAR) left (256 - input_size) // 2 top (256 - input_size) // 2 img img.crop((left, top, left input_size, top input_size)) arr np.asarray(img, dtypenp.float32) / 255.0 arr np.transpose(arr, (2, 0, 1)) arr (arr - MEAN) / STD arr np.expand_dims(arr, axis0).astype(np.float32) return arr def softmax(logits): logits logits - np.max(logits, axis-1, keepdimsTrue) exp np.exp(logits) return exp / exp.sum(axis-1, keepdimsTrue) def infer(session, image_path): input_name session.get_inputs()[0].name output_name session.get_outputs()[0].name x load_and_preprocess(image_path) logits session.run([output_name], {input_name: x})[0] probs softmax(logits)[0] topk_idx np.argsort(probs)[::-1][:5] topk_probs probs[topk_idx] return topk_idx.tolist(), topk_probs.tolist() session ort.InferenceSession(vit_tiny_sim.onnx, providers[CPUExecutionProvider]) idx, probs infer(session, test.jpg) print(idx, probs)这份代码里唯一需要根据实际需求调整的就是预处理里的Resize和Crop大小以及均值方差是否与训练配置一致。只要配置文件里的预处理顺序没变这套代码就能复用。5.3 性能分析和线程数调优ViT模型推理比同等规模的CNN耗时更大因为注意力机制的矩阵乘法计算密集而且ONNX Runtime在CPU上会把注意力里的多个算子拆开逐段执行。一个DeiT-Tiny模型在普通CPU上4线程的情况下单张224x224图片实测大概在20到30毫秒之间如果是ViT-Base可能要100毫秒以上。如果对延迟不满意一定要先做Profiling再想办法优化。建议用小工具单独统计预处理耗时、session.run耗时、后处理耗时。真实生产环境里我见过很多团队把session.run优化到60毫秒结果预处理一个Resize反而耗了80毫秒最后整体延迟依然很高。直接用numpy和PIL实现的预处理在CPU上其实不算快如果QPS要求高可以换用opencv的INTER_LINEAR或者用albumentations库做批处理。线程数调优可以用一段简单的扫描脚本import time import onnxruntime as ort import numpy as np x np.random.randn(1, 3, 224, 224).astype(np.float32) for threads in [1, 2, 4, 6, 8, 12]: opts ort.SessionOptions() opts.intra_op_num_threads threads sess ort.InferenceSession(vit_tiny_sim.onnx, sess_optionsopts, providers[CPUExecutionProvider]) # warmup for _ in range(10): sess.run(None, {input: x}) times [] for _ in range(50): t0 time.perf_counter() sess.run(None, {input: x}) times.append(time.perf_counter() - t0) print(threads, round(np.mean(times) * 1000, 2), ms)热身之后再计时非常重要因为onnxruntime第一次推理会触发一些初始化和算子融合动作如果不做热身数据会偏大很多。扫描结果能直观看到线程数从1到8延迟下降到12可能不再明显甚至变差这就是该机器的最优线程数。6. 常见问题排查与部署心得6.1 问题速查表整理一下这次过程中遇到的高频问题虽然不保证覆盖所有环境但大概率能帮到你问题现象可能原因解决办法导出时提示算子不支持自定义模块或新opset算子改成固定Shape降低opset或升级onnxruntimeonnxruntime加载时报错onnxruntime版本与导出的opset不匹配升级onnxruntime到与opset匹配的版本输出shape和类别数对不上模型head配置与预期不符检查训练配置里的num_classes不要硬编码CPU推理特别慢没有做onnxsim简化线程数不合适onnxsim简化模型扫描线程数GPU推理比CPU还慢数据量小或CUDA Provider未生效检查providers顺序与onnxruntime-gpu安装精度和PyTorch对不上预处理通道顺序、归一化或Resize方式不一致严格复现mmcls的pipeline用同一张图对比输出动态分辨率推理报错位置编码维度不匹配固定224x224导出或用插值方式重建位置编码这份表是我实际排查顺序的浓缩。绝大多数模型部署问题最终都能归因到“算子兼容性”和“数据一致性”两类其他都是这两类的变体。6.2 一个让我印象深刻的精度事故有一次我拿一个OpenMMLab导出的ViT模型在不同机器上跑离线评测结果有一台机器的top-1精度掉了3个点当时第一反应是onnxruntime版本问题。后来抓出输入数据一对比才发现是那台机器上安装的Pillow版本比较老对Image.resize的方法解释和训练环境不一致导致预处理产出完全不同。问题不在模型也不在ONNX Runtime而在图像库版本差异。这件事之后我养成了一个习惯部署包里直接固定Pillow、numpy、onnxruntime的版本号并且在推理代码里加入一个调试开关可以输出预处理后的样本均值、方差和PyTorch侧输出做一个简单的hash校验。数据一致性得到保证后才能继续调模型和算子。另外在前后端分离的团队里模型推理代码和训练代码很可能是不同人维护的最好在模型配置或者README里把预处理参数写清楚免得别人接手时靠猜。还有一个实用技巧导出ONNX后不要急着把PyTorch环境卸载。建议保留一份带mmcls库的虚拟环境专门用来做“基准校验”。每次改完推理代码或换了模型就跑一遍同一张测试图比较PyTorch logits和onnxruntime logits的差异。如果最大误差超过1e-3说明中间环节出问题了如果在1e-4以内说明整个链路是可靠的。最后再说一个经验ViT模型部署时不要一上来就追求全动态Shape。虽然技术方向上动态Shape是通用方案但落到VisionTransformer上位置编码、注意力矩阵的复杂度都跟输入尺寸强相关动态输入会让性能优化很难做也更容易在特定尺寸下触发算子兼容问题。绝大多数实际业务场景里图像分类的输入尺寸是固定的只要batch维度预留动态就足够满足高并发和多batch推理。先固定尺寸拿最低延迟再根据业务需要逐步开放动态能力这条路走起来会稳得多。