1. 项目概述:这条链路为什么值得折腾先聊点实际的。熟悉深度学习部署的朋友应该都有这种感觉:训练一套模型是一回事,把它高效地跑在目标环境里是另一回事。模型训练完,一堆评测指标都很好,到了真正要落地的时候,各种麻烦才冒出来——推理框架选什么、模型格式怎么转、动态输入怎么处理、精度掉了多少、显存占用怎么控制。这个项目要解决的问题,就是一条非常典型的部署链路:用 OpenMMLab 训练/导出 Vision Transformer 模型,转成 ONNX 格式,再用 onnxruntime 做推理。OpenMMLab 是我们在国产深度学习生态里绕不开的一座大山。它旗下有众多分类、检测、分割工具箱,学术界刷榜单、工业界做落地的都在用。Vision Transformer(简称 ViT)是近年来把 Transformer 架构从自然语言处理搬到计算机视觉领域的最有影响力的模型之一。至于 onnxruntime,微软出品的跨平台推理引擎,性能、部署灵活性、生态支持都在线。把它们拼在一起,解决了什么样的问题?最直观的痛点有三个:第一,训练框架和推理框架解耦。PyTorch 虽然训练方便,但直接部署对目标环境的要求太苛刻,显存和依赖都放不下。ONNX 作为一种开放的模型中间表示,可以做到“一次导出,多处部署”。第二,为模型换一个更轻量、可控的推理后端。onnxruntime 在 CPU 和 GPU 上都做了大量优化,很多场景下比直接用 PyTorch 跑要快,而且部署包体积小得多。第三,为跨平台移植铺路。ONNX Runtime 提供 Python、C++、C#、Java 等语言接口,配合 ONNX 格式,可以在边缘设备、服务器、移动端之间无缝迁移。这项目适合谁?如果你已经在用 OpenMMLab 训练模型,但还停留在“训练完就存个权重文件”的阶段,或者你被领导要求在服务器上把 ViT 模型跑起来供业务调用,再或者你想彻底搞懂从 PyTorch 到 ONNX 再到推理引擎这条标准部署流水线,这篇文章就是为你准备的。整条链路我实际操作下来,绕了几个不大不小的弯,把真正有价值的东西沉淀下来,按“为什么这么做、怎么做得对、踩过哪些坑”的顺序展开。整体路线如下:先用 OpenMMLab 生态中的分类工具箱训练一个 ViT 模型,然后通过官方部署工具链把权重文件转成 ONNX,最后分别用 onnxruntime 的 Python 接口和 C API 完成两套推理实现,做精度比对和性能调优。2. 整体设计与技术选型:为什么是 OpenMMLab + ViT + ONNX Runtime2.1 模型转换方案选型:官方工具链优先做模型部署,第一个问题就是:模型怎么从训练框架里弄出来?方案其实不少,我逐一分析过。最省事的思路是直接用 PyTorch 的torch.onnx.export,把训练好的模型权重读进来导出 ONNX。这个方法直观,但有几个隐患。OpenMMLab 的模型顶层封装非常复杂,包含各种预处理、后处理、数据并行包装逻辑。直接调用torch.onnx.export时,稍不留神就把一堆用不到的运算也导进去了,生成的 ONNX 图冗余节点多,推理性能难看,还容易碰到“动态 shape 导不出来”“某个算子不支持”这种报错。另一个思路是手工把权重迁移到一份干净的 ViT 实现里去,再用标准接口导出。工作量大,而且容易在权重拷贝过程中出错。最靠谱的方案是使用 OpenMMLab 官方提供的模型部署工具链。我这次用的是mmdeploy,它专门干这件事:把 OpenMMLab 系列的模型无缝转换为 ONNX、TensorRT、OpenVINO 等多种后端格式。好处在于,它把数据预处理(Normalize、Resize)、模型 forward、后处理(Softmax、TopK)整个计算图完整地导出,推理时的行为与训练时保持一致。至于推理引擎,我选择 onnxruntime 而不是 TensorRT,原因是项目目标环境没有专用显卡,必须在 CPU 上做部署,而 ONNX Runtime 的 CPU 性能和兼容性都很不错。2.2 Vision Transformer 导出时和 CNN 有什么本质区别选 ViT 还有一个值得说道的点:它和传统 CNN 在部署导出时遇到的挑战完全不同。CNN 模型(比如 ResNet)的核心算子相对固定,卷积、池化、标准化、激活函数,这些都是 ONNX 里一等公民。而 ViT 这类 Transformer 模型的核心是多头自注意力机制(Multi-Head Self-Attention),里面包含大量的reshape、transpose、matmul、softmax操作。这就好比搭积木。CNN 的积木块都是标准件,说明书上都有。Transformer 的积木块虽然也是那几个基础算子,但组合方式复杂,尺寸还经常变化。比如多头注意力中把(B, N, D)变成(B, num_heads, N, head_dim),一个transpose写不好,后面节点算出的张量维度就全对不上。更麻烦的是,ViT 的forward函数内部大量使用动态 shape。输入图片尺寸一旦不确定,序列长度N就跟着变,所有矩阵乘法的维数也跟着变。这要求导出时必须妥善处理动态轴,否则模型就只能跑固定尺寸的输入。还有一个隐藏得很深的坑:PyTorch 里的很多算子组合,在 ONNX 导出时匹配不到标准算子,会退化成子图(subgraph)。一旦 ONNX Runtime 找不到对应的高效实现,推理速度会非常感人。2.3 转换与推理链路总览整体链路可以用一句话概括:PyTorch 训练 → TorchScript 中间表示 → ONNX 计算图 → ONNX Runtime 推理引擎。为什么需要 TorchScript 中间表示这一步?因为 OpenMMLab 的工具链内部先把 PyTorch 模型跟踪为 TorchScript,再由 TorchScript 翻译成 ONNX。这带来一个额外好处:TorchScript 对动态控制流的表达能力更强,某些 PyTorch 原生导不出的结构可以绕过去。链路如下图所示:PyTorch 权重 (.pth) │ │ mmdeploy 的 torch2onnx 转换 ▼ ONNX 模型文件 (.onnx) │ │ 加载运行时 + 输入预处理 ▼ ONNX Runtime (CPU/GPU) │ ▼ 推理结果(类别 + 置信度)3. 环境准备与模型导出:实际操作记录3.1 环境配置与依赖安装先交代我用的环境,方便你对照排查。这里有个原则:版本对齐比什么都重要。OpenMMLab 全家桶对版本极敏感,PyTorch、CUDA、各工具包的版本错一个号都可能出兼容性问题。我推荐用 conda 建独立环境,不要污染系统 Python:conda create -n openmmlab python=3.8 -y conda activate openmmlab # 安装 PyTorch,这里是 CPU 版即可(部署目标也是 CPU) pip install torch==1.13.1 torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cpu # 安装 OpenMMLab 系列核心组件 pip install -U openmim mim install mmcv-full==1.7.1 mim install mmcv mim install mmcls # 安装部署工具链和推理运行时 pip install mmdeploy pip install onnx==1.13.1 pip install onnxruntime==1.14.1提示:mmdeploy 的安装分为两部分,基础包和对应后端的扩展包。这里只要 ONNX Runtime 后端,不需要装 TensorRT 扩展。装完之后,立刻跑一个最小测试,确认几个关键库能互相导入:python -c "import torch, mmcv, mmcls, mmdeploy, onnx, onnxruntime; print('all ok')"报错“No module named mmcls”之类,先检查是不是装成了mmclassification旧版本。新版统一改名mmcls,老教程里经常混着写,别被带偏。3.2 导出 ONNX 的完整命令与参数解析模型准备阶段,我用 mmcls 训练了一个以 ViT-Base 为骨干的分类模型,输入尺寸 224x224,类别数 1000。训练完事得到的权重文件是vit_base_epoch_100.pth。导出步骤其实就一条命令,关键是参数。python tools/deploy.py \ configs/mmcls/classification_onnxruntime_static.py \ /path/to/mmcls/configs/vit_base_224.py \ vit_base_epoch_100.pth \ /path/to/mmcls/demo/demo.jpeg \ --work-dir ./work_deploy \ --device cpu \ --dump-info逐个拆解参数含义:参数含义我的建议第一个 configmmdeploy 提供的部署配置模板,定义了导出格式、输入尺寸、动态轴策略静态图就选..._static.py,动态图选..._dynamic.py第二个 configmmcls 的模型定义文件必须与训练时一致,不能随意换权重文件模型的.pth参数文件确保和模型定义能对上demo 图片转换时用来跑一次 forward,用于追踪计算图随便一张正常图片即可--device cpu导出图和做验证推理用的设备如果模型太大可以用 cuda,但导出结果一致--dump-info同时输出推理所需的 meta 信息建议默认开启,后面写 C++ 接口用得上导出完成后,work_deploy目录下会有几个关键文件:end2end.onnx:完整可推理的 ONNX 模型end2end.engine:如果是 TensorRT 后端才会有,ONNX 后端没有pipeline.json:记录模型输入输出的名字、shape、预处理参数这一步也是坑最多的地方。我建议你导出后立刻做两件事:第一,用onnx.checker.check_model做合法性检查;第二,用onnxruntime跑一次 dummy 输入,保证前向能通。3.3 用 MIM 快速获取预训练模型的另一个思路自己训练虽然对理解模型有好处,但在探索部署链路时,为了节约时间,建议直接用 MIM 从 OpenMMLab 模型库下载官方预训练权重。这一招很实用:mim download mmcls --config vit-base-p16_224-80e_imagenet --dest ./models执行完,目录里会出现对应的config文件和pth权重文件。省去训练的漫长等待,直接把精力集中在部署环节。我后来调试 C API 时就是这么干的,效果很好。下载的官方权重如果和本地训练权重名词不同,在部署配置里把load_from字段指到新权重即可。4. onnxruntime 推理实现:Python 与 C API 两套方案4.1 Python 推理脚本:快速打通链路导出的 ONNX 模型如何跑起来?先上 Python 脚本,核心思路就三步:加载模型、构造输入、跑run方法。import onnxruntime as ort import numpy as np import cv2 from mmcv.imgs import imread, imrescale, imnormalize def preprocess(image_path, input_size=224, mean=(123.675, 116.28, 103.53), std=(58.395, 57.12, 57.375)): img = imread(image_path, channel_order='bgr') img = imrescale(img, (input_size, input_size)) img = imnormalize(img, mean, std, to_rgb=True) img = img.transpose(2, 0, 1)[None, ...].astype(np.float32) return img def infer_onnx(model_path, image_path): sess = ort.InferenceSession(model_path, providers=['CPUExecutionProvider']) input_name = sess.get_inputs()[0].name output_name = sess.get_outputs()[0].name x = preprocess(image_path) outputs = sess.run([output_name], {input_name: x}) probs = np.squeeze(outputs[0]) top1_idx = int(np.argmax(probs)) return top1_idx, float(probs[top1_idx]) if __name__ == '__main__': idx, score = infer_onnx('./work_deploy/end2end.onnx', 'test.jpg') print(f'top1 class index: {idx}, score: {score:.4f}')关于预处理,需要特别强调一点:它的顺序、参数、色彩空间转换必须和训练时候完全一致。我一开始就是从这儿栽的跟头,图片预处理差了一个to_rgb,推理出来的类别直接错到十万八千里。OpenMMLab 内部默认图像通道顺序是 BGR,归一化的均值方差也是固定的,千万别自作主张改成skimage或者PIL的标准流程。跑完脚本,如果输出的类别索引和用 mmcls 原模型推理的结果一致,恭喜你,整条链路已经打通。这一步顺利的话,实际耗时在几百毫秒到一两秒之间,取决于输入尺寸和 CPU 性能。4.2 数据格式与张量排布细节ONNX Runtime 的输入输出都是numpy.ndarray,这和大家用 PyTorch 时候的感受完全不同。PyTorch 里张量默认NCHW排布,numpy 数组默认是NHWC排布(其实没有严格的维度语义,只看你怎么解释)。转换时必须显式做成(1, 3, 224, 224)的形状,否则模型推理会直接报维度不匹配的错误。还有一个非常隐蔽的坑:图像数据格式要连续内存块。如果你对 numpy 数组做切片或者做np.expand_dims,得到的可能是非连续内存,虽然run方法一般会自动处理,但性能会打折扣。稳妥的做法是在输入前加一行np.ascontiguousarray。x = np.ascontiguousarray(x, dtype=np.float32)4.3 C API 推理:服务化部署的进阶之路Python 脚本验证后,如果你的模型最终要集成到 C++ 服务里,需要换用 ONNX Runtime 的 C API。这一步比 Python 麻烦一些,但也更能锻炼对内存和生命周期的理解。核心代码骨架:#include <onnxruntime_cxx_api.h> #include <vector> #include <iostream> int main() { Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "vit-inference"); Ort::SessionOptions session_options; session_options.SetIntraOpNumThreads(4); session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); // 模型路径用宽字符的 wstring,这个细节很容易忽略 auto model_path = std::wstring(L"./work_deploy/end2end.onnx"); Ort::Session session(env, model_path.c_str(), session_options); // 从 session 读取输入输出信息 Ort::AllocatorWithDefaultOptions allocator; auto input_name = session.GetInputNameAllocated(0, allocator); auto input_shape = session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape(); auto output_name = session.GetOutputNameAllocated(0, allocator); // 构造输入张量 std::vector<int64_t> input_shape_vec = {1, 3, 224, 224}; std::vector<float> input_data(1 * 3 * 224 * 224); // ... 这里把图像数据经过预处理后填入 input_data Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_data.data(), input_data.size(), input_shape_vec.data(), input_shape_vec.size()); // 推理 auto output_tensors = session.Run(Ort::RunOptions{nullptr}, &input_name, &input_tensor, 1, &output_name, 1); // 读取输出 auto output_info = output_tensors.front().GetTensorTypeAndShapeInfo(); size_t output_size = output_info.GetElementCount(); float* output_data = output_tensors.front().GetTensorMutableData<float>(); for (size_t i = 0; i < 5; i++) { std::cout << "class " << i << " score: " << output_data[i] << std::endl; } return 0; }编译时别忘了链接 ONNX Runtime 库:g++ -std=c++17 infer.cpp -o infer \ -I /path/to/onnxruntime/include \ -L /path/to/onnxruntime/lib \ -lonnxruntime需要设置动态库路径才能运行:export LD_LIBRARY_PATH=/path/to/onnxruntime/lib:$LD_LIBRARY_PATH这一段工程实现有个好处:它可以直接嵌入现有的 C++ 推理服务,不用像 Python 那样必须拉起一个解释器进程,内存和启动时间都节省了不少。我在实际项目里就是把这段逻辑封装成一个VitInference类,只暴露LoadModel和Predict两个接口,其他细节全藏起来。5. 性能对比与优化:实际数据说话5.1 三套方案的实测对比性能测试是部署环节的重头戏。我在一台 Intel Xeon Silver 系列 CPU(8 核、无 GPU)的服务器上,分别测试了三套推理方案:PyTorch 直接推理(训练时的标准 forward)ONNX Runtime Python 接口ONNX Runtime C API 接口统一使用 224x224 输入、批量大小为 1、纯 CPU 推理,数据如下:推理方案单张耗时(ms)峰值内存占用(MB)依赖体积PyTorch(CPU)约 750约 1200巨大(需完整框架)onnxruntime Python约 420约 350中等onnxruntime C API约 350约 280小(仅一个动态库)PyTorch 和 ONNX Runtime 的差距不是一点半点。这里面的原因不复杂:ONNX Runtime 在计算图层面做了融合优化(例如把 BN 和卷积融合成一个节点、把连续 reshape 合并)、线程池调度更精细、还有专门针对 Transformer 子图的优化规则。CPU 环境下,这些优化直接反映在推理延迟上。5.2 影响推理速度的六个因素如果你复测时发现性能不理想,可以用控制变量法逐项排查。我的经验是关注这六个变量:线程数设置。ONNX Runtime 默认会用满所有核心,但有时调成SetIntraOpNumThreads(4)比默认反而更快,因为少了线程切换的开销,尤其在线程数量超过物理核心数时更明显。图优化级别。ORT_ENABLE_ALL级优化会尝试各种计算融合,务必开启。它有时候能带来接近一倍的速度提升。输入数据连续性与对齐。非连续内存的输入会迫使运行时做一次额外拷贝。量化与否。如果精度容忍度高,可以把 FP32 模型转成 INT8 量化模型,CPU 推理速度能再上一个量级。ONNX Runtime 支持动态量化和静态量化两条路线,静态量化精度更高。动态轴是否避免了不必要的形状推断。如果模型里大量出现动态 shape,运行时会在每次推理时重新做形状推导,拖慢速度。固定输入尺寸的静态图是性能最好的形态。注意力模块实现方式。ViT 里最耗时的就是多头注意力。某些 PyTorch 实现用了einops、flash-attn等高级库,导出 ONNX 时代码结构不好,性能直接崩盘。速度异常时用图优化工具看结构是重要的排查手段。有一个很典型的对比案例:同一套 ViT-Base,把头数num_heads从 12 改成 8,序列长度从 196 改成 64,ONNX 导出后的计算图节点数相差三成。这说明模型设计阶段的超参选择对部署性能有直接影响。5.3 用 onnxruntime 自带的 profiling 工具定位瓶颈如果你还是不满足,或者碰到诡异的性能退化,一定要会用 onnxruntime 自带的分析工具。开启 profiling 的方法很简单:so = ort.SessionOptions() so.enable_profiling = True sess = ort.InferenceSession('end2end.onnx', so) # 跑一次推理 sess.run(...) # 结束后会生成一个 .json 文件接着用 Chrome 的chrome://tracing加载这个 JSON 文件,就能看到每个算子消耗的时间和执行顺序。哪一层慢、哪里多了一堆奇怪的Cast节点、哪个Transpose偷走了大量时间,全部一目了然。学会这一手,比盲目调参有效十倍。6. 常见问题与排查方法:我踩过的十个坑,一次说清从 PyTorch 模型到 ONNX Runtime 推理,出错的地方不少。按出现频率排序,把最常见的坑和排查步骤整理如下:6.1 模型导出的常见报错报错现象根因分析解决方案Unsupported operator aten::xxxPyTorch 实现里的某个算子 ONNX 导出器不认识更新 PyTorch 和 mmdeploy 版本;手动导出前把该算子替换为等价算子组合;检查是不是需要更新到更新的 ONNX opsetRuntimeError: Failed to export且错误指向reshapeViT 中动态 shape 导致 reshape 目标维度无法推断设置动态轴(dynamic_axes);用固定输入尺寸;替换内部实现导出成功但推理结果全错数据预处理顺序不对;模型加载前忘记eval模式逐项对比预处理细节,特别是通道顺序和归一化参数;导出前声明model.eval()静态图导出提示None维度数据模型里有未初始化的张量,通常是 batch size 引起的用 dummy input 跑一次 forward 后再导出,而非直接传 NoneGraph must be in single static assignment (SSA) form模型图中存在变量先赋值后重新赋值的张量查看是否有循环内的张量复用;用 trace 而非 script6.2 onnxruntime 推理端常见问题报错现象根因分析解决方案Invalid Symbolic Shape模型含动态轴,session 初始化时无法确定输出形状推理前固定输入 shape;或改用动态 shape 配置;或设置session_options.SetDynamicShape相关选项推理时出现Segmentation fault常见于 C API 中输入张量大小与实际数据不同;内存未对齐;tensor 生命周期提前释放检查CreateTensor时传入的数据指针和大小是否一致;确保input_data在Run前一直存活No implementation found for ...图里包含后端不支持的算子,或后端版本太旧升级 onnxruntime;换用支持该算子的执行提供器(如 TensorRT)或等待算子实现升级Python 接口运行很慢,C API 不慢Python 调用 overhead 在批量小的时候体现明显使用IOBinding减少 Python 层 GIL 开销;增加 batch;把多次推理合并为一个 request图像预处理后输入值范围不对模型期望 0~255 范围的图像输入,你归一化到了 0~1;或反之检查训练时的图像标准化方式,从配置文件的 pipeline 中读取原始值,严格保持一致6.3 一个必须强调的老生常谈:版本一致性很多问题最终都能归结到版本不匹配。onnxruntime 里的算子实现即使算子名相同,不同版本的行为也差异很大。ONNX 的 opset 从 11 到 20,版本越高,支持的算子越丰富,但对老库的兼容性反而可能出问题。如果你不是必须用最新特性,我建议锁死一个稳定组合:PyTorch 1.13 + ONNX 1.13 + onnxruntime 1.14,这组配对我实测下来兼容性最佳。如果手头环境已经换代,至少保证 ONNX 模型的 opset 版本不能高于 onnxruntime 支持的版本。一旦出现算子解析失败,优先检查这个不变量。7. 给 ONNX 模型加个加速 Buff:ONNX Runtime 的线程与内存配置上面提到性能优化时,线程和内存配置值得单独拎出来展开。很多人图省事,直接把默认配置跑完就结束了,结果服务端并发上来之后延迟翻了倍,这时才发现运行时配置有多重要。7.1 SessionOptions 里几个关键开关ONNX Runtime 支持一个专门的配置结构体OrtSessionOptions,其中的参数决定运行时如何调度计算图。SetIntraOpNumThreads:控制单个算子内部的并行线程数。矩阵乘法这类算子内部并行度高,适用于多线程,但线程太多会引入开销。经验值:设置为物理核心数的一半到三分之二。SetInterOpNumThreads:控制计算图中不同算子之间的并行执行。如果你的模型算子是串行依赖关系强,这个参数作用不大;如果模型图结构中有明显的并行分支,把它调大就有收益。SetGraphOptimizationLevel:三档可选,从ORT_DISABLE_ALL到ORT_ENABLE_ALL,直接决定运行时是否对计算图做算子融合。务必设置为ORT_ENABLE_ALL。还有一个容易忽略的是内存分配器。默认情况下 onnxruntime 使用 arena 分配器,适合长生命周期场景。但如果你的推理进程会频繁创建和销毁 session,arena 的内存池反而会造成内存不回收的假象,看起来像泄漏。这时可以设置session_options.AddConfigEntry("session.use_arena", "0")关掉 arena。7.2 遇到批量推理时该如何调整如果你的业务不是一张一张图推理,而是需要一次处理多张,有两种姿势可选。第一种是普通 batch 推理,直接把输入做成(N, 3, H, W)的 N 张图一起送入。这样做吞吐率高,但延迟可能被最慢的那张图拖累。第二种是并发推理,在同一 session 上开多个线程各自跑不同输入。ONNX Runtime 的 session 是线程安全的,多个线程同时调用Run不会冲突。这在 Python 接口上,可以结合concurrent.futures.ThreadPoolExecutor用;在 C API 下,只需各自管理自己的Ort::Value。实测下来,并发推理比单纯增 batch 更灵活,适合服务化场景。8. 工程化落地的几个补充细节8.1 用 onnx 工具库做模型体检模型导出来后,在交给 onnxruntime 之前,建议先做一次全面的模型体检。我在导出完成后用各工具跑了一遍,意义很大。import onnx from onnx import checker, helper, shape_inference model = onnx.load("./work_deploy/end2end.onnx") checker.check_model(model, full_check=True) # 做一次形状推断,能暴露隐藏的维度错误 inferred_model = shape_inference.infer_shapes(model) onnx.save(inferred_model, "./work_deploy/end2end_inferred.onnx") # 统计节点类型,看是否有意外出现的算子 from collections import Counter node_types = Counter(n.op_type for n in model.graph.node) print(node_types.most_common(20))最后这行统计很有用。我遇到过一种情况,模型里出现了大量的Cast节点,原因是 ViT 的某些算子对精度有要求,被强制转换成了double。这种精度提升对分类任务没有明显收益,却让 CPU 推理时间翻倍。发现之后,我修改了导出配置,统一用 FP32,速度立刻就回来了。8.2 模型输入的 meta 信息管理服务化部署时,把模型文件和一些元信息放在一起管理很重要。模型除了本身的 ONNX 图,还有输入输出节点名称、图像尺寸、归一化参数等关联信息。mmdeploy 的--dump-info参数会生成pipeline.json,里面包含了上述信息,但格式不直观。我的做法是自己在模型旁边单独放一个 JSON 描述文件:{ "model": "end2end.onnx", "input": { "name": "input", "shape": [1, 3, 224, 224], "dtype": "float32" }, "output": { "name": "output", "shape": [1, 1000], "dtype": "float32" }, "preprocess": { "mean": [123.675, 116.28, 103.53], "std": [58.395, 57.12, 57.375], "color_order": "BGR", "resize": "224x224" } }有了这个文件,模型文件本身就可以做到平台无关。换任何语言、任何框架来接 ONNX 模型时,参数的读取都有了统一来源,不会出现模型文件换台机器就跑不起来的窘境。8.3 单元测试与精度比对方案模型的精度验证是部署里最容易马虎又最不该马虎的环节。我的精度验证标准方案是:选 100 张测试图,分别用 PyTorch 原模型和 ONNX Runtime 推理。对每张图比较两边的 top-1 类别是否一致。对置信度分数,计算平均绝对误差(MAE)。误差通常在1e-4到1e-5量级,这是 FP32 精度在不同算子实现下的正常波动。如果误差超过1e-2,基本可以断定是预处理、数据格式或算子精度出了问题。这个小脚本虽然简单,但有效。我记得有一次导出的模型在测试集上准确率掉了 3 个百分点,排查了整整半天,最后发现图像归一化的std写成了训练时使用的mean,类似这种错误在视图比对时能很快暴露问题。9. 再延伸一步:为 ONNX 推理加个“加速 Buff”,动态输入与量化探索如果前面都没问题,还有两个方向可以继续挖掘,把 ViT 模型的部署性能再往前推一点。方向一:动态输入尺寸的支持。ViT 在导出时如果用动态轴配置,就可以输入任意尺寸的图片,而不必强制 reshape 到 224x224。这对真实场景很有意义,客户发来的图片构不成正方形,直接捏成 224x224 会损失大量信息。动态轴的配置在 mmdeploy 的配置文件中可以指定:dynamic_axes = { "input": {0: "batch", 2: "height", 3: "width"}, "output": {0: "batch"} }配置之后重新导出,就能支持任意尺寸输入。代价是推理时会有额外的 shape 推断开销,每次输入尺寸如果和上次不同,ONNX Runtime 需要重新编译部分算子,延迟会略有上升。如果你场景里图片尺寸变化不频繁,动态输入带来的精度收益远大于这点开销。方向二:量化。CPU 部署模型时,INT8 量化是提升速度最有效的手段。ONNX Runtime 提供onnxruntime.quantization工具包,静态量化需要一批校准数据来计算激活值的动态范围:from onnxruntime.quantization import quantize_static, QuantType, CalibrationDataReader class DataReader(CalibrationDataReader): def __init__(self, data_paths): # 实现 get_next 方法,返回一批校准样本 pass quantize_static( model_input="end2end.onnx", model_output="end2end_int8.onnx", calibration_data_reader=DataReader([...]), quant_format=QuantType.QInt8 )量化后的 ViT 模型在 CPU 上通常能获得 2~3 倍的速度提升,精度损失在分类任务中一般在 1%~2% 以内。这条路线如果业务需求有压测要求,非常值得深入研究。我的总结体会这条链路跑通之后,往后遇到任何 OpenMMLab 系模型的部署需求,基本都能沿着一套模式完成:模型准备、导出 ONNX、选择推理后端、处理输入输出、性能调优。整个过程看下来,真正费时间的从来不是写代码,而是版本匹配、预处理一致性、动态 shape 这些问题。一个朴素的想法是,跨框架部署的本质不是“模型格式转换”,而是把训练时的所有假设——包括数据的进出、算子的边界、精度容差——原原本本地翻译到另一端。只要抓住了这一条,很多诡异问题其实都能一眼看穿。最后再分享一个我个人的小习惯:每导出一个新模型,我都会把环境依赖的版本号写进一个 requirements 文件里,和模型文件一起保存。三个月后再回来看,依赖已经大变样了,但只要有那份文件,任何方法都能快速复现当时的部署环境。这个习惯帮我避过太多次的版本灾难,建议你从现在开始也这样做。
1. 项目概述:这条链路为什么值得折腾先聊点实际的。熟悉深度学习部署的朋友应该都有这种感觉:训练一套模型是一回事,把它高效地跑在目标环境里是另一回事。模型训练完,一堆评测指标都很好,到了真正要落地的时候,各种麻烦才冒出来——推理框架选什么、模型格式怎么转、动态输入怎么处理、精度掉了多少、显存占用怎么控制。这个项目要解决的问题,就是一条非常典型的部署链路:用 OpenMMLab 训练/导出 Vision Transformer 模型,转成 ONNX 格式,再用 onnxruntime 做推理。OpenMMLab 是我们在国产深度学习生态里绕不开的一座大山。它旗下有众多分类、检测、分割工具箱,学术界刷榜单、工业界做落地的都在用。Vision Transformer(简称 ViT)是近年来把 Transformer 架构从自然语言处理搬到计算机视觉领域的最有影响力的模型之一。至于 onnxruntime,微软出品的跨平台推理引擎,性能、部署灵活性、生态支持都在线。把它们拼在一起,解决了什么样的问题?最直观的痛点有三个:第一,训练框架和推理框架解耦。PyTorch 虽然训练方便,但直接部署对目标环境的要求太苛刻,显存和依赖都放不下。ONNX 作为一种开放的模型中间表示,可以做到“一次导出,多处部署”。第二,为模型换一个更轻量、可控的推理后端。onnxruntime 在 CPU 和 GPU 上都做了大量优化,很多场景下比直接用 PyTorch 跑要快,而且部署包体积小得多。第三,为跨平台移植铺路。ONNX Runtime 提供 Python、C++、C#、Java 等语言接口,配合 ONNX 格式,可以在边缘设备、服务器、移动端之间无缝迁移。这项目适合谁?如果你已经在用 OpenMMLab 训练模型,但还停留在“训练完就存个权重文件”的阶段,或者你被领导要求在服务器上把 ViT 模型跑起来供业务调用,再或者你想彻底搞懂从 PyTorch 到 ONNX 再到推理引擎这条标准部署流水线,这篇文章就是为你准备的。整条链路我实际操作下来,绕了几个不大不小的弯,把真正有价值的东西沉淀下来,按“为什么这么做、怎么做得对、踩过哪些坑”的顺序展开。整体路线如下:先用 OpenMMLab 生态中的分类工具箱训练一个 ViT 模型,然后通过官方部署工具链把权重文件转成 ONNX,最后分别用 onnxruntime 的 Python 接口和 C API 完成两套推理实现,做精度比对和性能调优。2. 整体设计与技术选型:为什么是 OpenMMLab + ViT + ONNX Runtime2.1 模型转换方案选型:官方工具链优先做模型部署,第一个问题就是:模型怎么从训练框架里弄出来?方案其实不少,我逐一分析过。最省事的思路是直接用 PyTorch 的torch.onnx.export,把训练好的模型权重读进来导出 ONNX。这个方法直观,但有几个隐患。OpenMMLab 的模型顶层封装非常复杂,包含各种预处理、后处理、数据并行包装逻辑。直接调用torch.onnx.export时,稍不留神就把一堆用不到的运算也导进去了,生成的 ONNX 图冗余节点多,推理性能难看,还容易碰到“动态 shape 导不出来”“某个算子不支持”这种报错。另一个思路是手工把权重迁移到一份干净的 ViT 实现里去,再用标准接口导出。工作量大,而且容易在权重拷贝过程中出错。最靠谱的方案是使用 OpenMMLab 官方提供的模型部署工具链。我这次用的是mmdeploy,它专门干这件事:把 OpenMMLab 系列的模型无缝转换为 ONNX、TensorRT、OpenVINO 等多种后端格式。好处在于,它把数据预处理(Normalize、Resize)、模型 forward、后处理(Softmax、TopK)整个计算图完整地导出,推理时的行为与训练时保持一致。至于推理引擎,我选择 onnxruntime 而不是 TensorRT,原因是项目目标环境没有专用显卡,必须在 CPU 上做部署,而 ONNX Runtime 的 CPU 性能和兼容性都很不错。2.2 Vision Transformer 导出时和 CNN 有什么本质区别选 ViT 还有一个值得说道的点:它和传统 CNN 在部署导出时遇到的挑战完全不同。CNN 模型(比如 ResNet)的核心算子相对固定,卷积、池化、标准化、激活函数,这些都是 ONNX 里一等公民。而 ViT 这类 Transformer 模型的核心是多头自注意力机制(Multi-Head Self-Attention),里面包含大量的reshape、transpose、matmul、softmax操作。这就好比搭积木。CNN 的积木块都是标准件,说明书上都有。Transformer 的积木块虽然也是那几个基础算子,但组合方式复杂,尺寸还经常变化。比如多头注意力中把(B, N, D)变成(B, num_heads, N, head_dim),一个transpose写不好,后面节点算出的张量维度就全对不上。更麻烦的是,ViT 的forward函数内部大量使用动态 shape。输入图片尺寸一旦不确定,序列长度N就跟着变,所有矩阵乘法的维数也跟着变。这要求导出时必须妥善处理动态轴,否则模型就只能跑固定尺寸的输入。还有一个隐藏得很深的坑:PyTorch 里的很多算子组合,在 ONNX 导出时匹配不到标准算子,会退化成子图(subgraph)。一旦 ONNX Runtime 找不到对应的高效实现,推理速度会非常感人。2.3 转换与推理链路总览整体链路可以用一句话概括:PyTorch 训练 → TorchScript 中间表示 → ONNX 计算图 → ONNX Runtime 推理引擎。为什么需要 TorchScript 中间表示这一步?因为 OpenMMLab 的工具链内部先把 PyTorch 模型跟踪为 TorchScript,再由 TorchScript 翻译成 ONNX。这带来一个额外好处:TorchScript 对动态控制流的表达能力更强,某些 PyTorch 原生导不出的结构可以绕过去。链路如下图所示:PyTorch 权重 (.pth) │ │ mmdeploy 的 torch2onnx 转换 ▼ ONNX 模型文件 (.onnx) │ │ 加载运行时 + 输入预处理 ▼ ONNX Runtime (CPU/GPU) │ ▼ 推理结果(类别 + 置信度)3. 环境准备与模型导出:实际操作记录3.1 环境配置与依赖安装先交代我用的环境,方便你对照排查。这里有个原则:版本对齐比什么都重要。OpenMMLab 全家桶对版本极敏感,PyTorch、CUDA、各工具包的版本错一个号都可能出兼容性问题。我推荐用 conda 建独立环境,不要污染系统 Python:conda create -n openmmlab python=3.8 -y conda activate openmmlab # 安装 PyTorch,这里是 CPU 版即可(部署目标也是 CPU) pip install torch==1.13.1 torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cpu # 安装 OpenMMLab 系列核心组件 pip install -U openmim mim install mmcv-full==1.7.1 mim install mmcv mim install mmcls # 安装部署工具链和推理运行时 pip install mmdeploy pip install onnx==1.13.1 pip install onnxruntime==1.14.1提示:mmdeploy 的安装分为两部分,基础包和对应后端的扩展包。这里只要 ONNX Runtime 后端,不需要装 TensorRT 扩展。装完之后,立刻跑一个最小测试,确认几个关键库能互相导入:python -c "import torch, mmcv, mmcls, mmdeploy, onnx, onnxruntime; print('all ok')"报错“No module named mmcls”之类,先检查是不是装成了mmclassification旧版本。新版统一改名mmcls,老教程里经常混着写,别被带偏。3.2 导出 ONNX 的完整命令与参数解析模型准备阶段,我用 mmcls 训练了一个以 ViT-Base 为骨干的分类模型,输入尺寸 224x224,类别数 1000。训练完事得到的权重文件是vit_base_epoch_100.pth。导出步骤其实就一条命令,关键是参数。python tools/deploy.py \ configs/mmcls/classification_onnxruntime_static.py \ /path/to/mmcls/configs/vit_base_224.py \ vit_base_epoch_100.pth \ /path/to/mmcls/demo/demo.jpeg \ --work-dir ./work_deploy \ --device cpu \ --dump-info逐个拆解参数含义:参数含义我的建议第一个 configmmdeploy 提供的部署配置模板,定义了导出格式、输入尺寸、动态轴策略静态图就选..._static.py,动态图选..._dynamic.py第二个 configmmcls 的模型定义文件必须与训练时一致,不能随意换权重文件模型的.pth参数文件确保和模型定义能对上demo 图片转换时用来跑一次 forward,用于追踪计算图随便一张正常图片即可--device cpu导出图和做验证推理用的设备如果模型太大可以用 cuda,但导出结果一致--dump-info同时输出推理所需的 meta 信息建议默认开启,后面写 C++ 接口用得上导出完成后,work_deploy目录下会有几个关键文件:end2end.onnx:完整可推理的 ONNX 模型end2end.engine:如果是 TensorRT 后端才会有,ONNX 后端没有pipeline.json:记录模型输入输出的名字、shape、预处理参数这一步也是坑最多的地方。我建议你导出后立刻做两件事:第一,用onnx.checker.check_model做合法性检查;第二,用onnxruntime跑一次 dummy 输入,保证前向能通。3.3 用 MIM 快速获取预训练模型的另一个思路自己训练虽然对理解模型有好处,但在探索部署链路时,为了节约时间,建议直接用 MIM 从 OpenMMLab 模型库下载官方预训练权重。这一招很实用:mim download mmcls --config vit-base-p16_224-80e_imagenet --dest ./models执行完,目录里会出现对应的config文件和pth权重文件。省去训练的漫长等待,直接把精力集中在部署环节。我后来调试 C API 时就是这么干的,效果很好。下载的官方权重如果和本地训练权重名词不同,在部署配置里把load_from字段指到新权重即可。4. onnxruntime 推理实现:Python 与 C API 两套方案4.1 Python 推理脚本:快速打通链路导出的 ONNX 模型如何跑起来?先上 Python 脚本,核心思路就三步:加载模型、构造输入、跑run方法。import onnxruntime as ort import numpy as np import cv2 from mmcv.imgs import imread, imrescale, imnormalize def preprocess(image_path, input_size=224, mean=(123.675, 116.28, 103.53), std=(58.395, 57.12, 57.375)): img = imread(image_path, channel_order='bgr') img = imrescale(img, (input_size, input_size)) img = imnormalize(img, mean, std, to_rgb=True) img = img.transpose(2, 0, 1)[None, ...].astype(np.float32) return img def infer_onnx(model_path, image_path): sess = ort.InferenceSession(model_path, providers=['CPUExecutionProvider']) input_name = sess.get_inputs()[0].name output_name = sess.get_outputs()[0].name x = preprocess(image_path) outputs = sess.run([output_name], {input_name: x}) probs = np.squeeze(outputs[0]) top1_idx = int(np.argmax(probs)) return top1_idx, float(probs[top1_idx]) if __name__ == '__main__': idx, score = infer_onnx('./work_deploy/end2end.onnx', 'test.jpg') print(f'top1 class index: {idx}, score: {score:.4f}')关于预处理,需要特别强调一点:它的顺序、参数、色彩空间转换必须和训练时候完全一致。我一开始就是从这儿栽的跟头,图片预处理差了一个to_rgb,推理出来的类别直接错到十万八千里。OpenMMLab 内部默认图像通道顺序是 BGR,归一化的均值方差也是固定的,千万别自作主张改成skimage或者PIL的标准流程。跑完脚本,如果输出的类别索引和用 mmcls 原模型推理的结果一致,恭喜你,整条链路已经打通。这一步顺利的话,实际耗时在几百毫秒到一两秒之间,取决于输入尺寸和 CPU 性能。4.2 数据格式与张量排布细节ONNX Runtime 的输入输出都是numpy.ndarray,这和大家用 PyTorch 时候的感受完全不同。PyTorch 里张量默认NCHW排布,numpy 数组默认是NHWC排布(其实没有严格的维度语义,只看你怎么解释)。转换时必须显式做成(1, 3, 224, 224)的形状,否则模型推理会直接报维度不匹配的错误。还有一个非常隐蔽的坑:图像数据格式要连续内存块。如果你对 numpy 数组做切片或者做np.expand_dims,得到的可能是非连续内存,虽然run方法一般会自动处理,但性能会打折扣。稳妥的做法是在输入前加一行np.ascontiguousarray。x = np.ascontiguousarray(x, dtype=np.float32)4.3 C API 推理:服务化部署的进阶之路Python 脚本验证后,如果你的模型最终要集成到 C++ 服务里,需要换用 ONNX Runtime 的 C API。这一步比 Python 麻烦一些,但也更能锻炼对内存和生命周期的理解。核心代码骨架:#include <onnxruntime_cxx_api.h> #include <vector> #include <iostream> int main() { Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "vit-inference"); Ort::SessionOptions session_options; session_options.SetIntraOpNumThreads(4); session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); // 模型路径用宽字符的 wstring,这个细节很容易忽略 auto model_path = std::wstring(L"./work_deploy/end2end.onnx"); Ort::Session session(env, model_path.c_str(), session_options); // 从 session 读取输入输出信息 Ort::AllocatorWithDefaultOptions allocator; auto input_name = session.GetInputNameAllocated(0, allocator); auto input_shape = session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape(); auto output_name = session.GetOutputNameAllocated(0, allocator); // 构造输入张量 std::vector<int64_t> input_shape_vec = {1, 3, 224, 224}; std::vector<float> input_data(1 * 3 * 224 * 224); // ... 这里把图像数据经过预处理后填入 input_data Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_data.data(), input_data.size(), input_shape_vec.data(), input_shape_vec.size()); // 推理 auto output_tensors = session.Run(Ort::RunOptions{nullptr}, &input_name, &input_tensor, 1, &output_name, 1); // 读取输出 auto output_info = output_tensors.front().GetTensorTypeAndShapeInfo(); size_t output_size = output_info.GetElementCount(); float* output_data = output_tensors.front().GetTensorMutableData<float>(); for (size_t i = 0; i < 5; i++) { std::cout << "class " << i << " score: " << output_data[i] << std::endl; } return 0; }编译时别忘了链接 ONNX Runtime 库:g++ -std=c++17 infer.cpp -o infer \ -I /path/to/onnxruntime/include \ -L /path/to/onnxruntime/lib \ -lonnxruntime需要设置动态库路径才能运行:export LD_LIBRARY_PATH=/path/to/onnxruntime/lib:$LD_LIBRARY_PATH这一段工程实现有个好处:它可以直接嵌入现有的 C++ 推理服务,不用像 Python 那样必须拉起一个解释器进程,内存和启动时间都节省了不少。我在实际项目里就是把这段逻辑封装成一个VitInference类,只暴露LoadModel和Predict两个接口,其他细节全藏起来。5. 性能对比与优化:实际数据说话5.1 三套方案的实测对比性能测试是部署环节的重头戏。我在一台 Intel Xeon Silver 系列 CPU(8 核、无 GPU)的服务器上,分别测试了三套推理方案:PyTorch 直接推理(训练时的标准 forward)ONNX Runtime Python 接口ONNX Runtime C API 接口统一使用 224x224 输入、批量大小为 1、纯 CPU 推理,数据如下:推理方案单张耗时(ms)峰值内存占用(MB)依赖体积PyTorch(CPU)约 750约 1200巨大(需完整框架)onnxruntime Python约 420约 350中等onnxruntime C API约 350约 280小(仅一个动态库)PyTorch 和 ONNX Runtime 的差距不是一点半点。这里面的原因不复杂:ONNX Runtime 在计算图层面做了融合优化(例如把 BN 和卷积融合成一个节点、把连续 reshape 合并)、线程池调度更精细、还有专门针对 Transformer 子图的优化规则。CPU 环境下,这些优化直接反映在推理延迟上。5.2 影响推理速度的六个因素如果你复测时发现性能不理想,可以用控制变量法逐项排查。我的经验是关注这六个变量:线程数设置。ONNX Runtime 默认会用满所有核心,但有时调成SetIntraOpNumThreads(4)比默认反而更快,因为少了线程切换的开销,尤其在线程数量超过物理核心数时更明显。图优化级别。ORT_ENABLE_ALL级优化会尝试各种计算融合,务必开启。它有时候能带来接近一倍的速度提升。输入数据连续性与对齐。非连续内存的输入会迫使运行时做一次额外拷贝。量化与否。如果精度容忍度高,可以把 FP32 模型转成 INT8 量化模型,CPU 推理速度能再上一个量级。ONNX Runtime 支持动态量化和静态量化两条路线,静态量化精度更高。动态轴是否避免了不必要的形状推断。如果模型里大量出现动态 shape,运行时会在每次推理时重新做形状推导,拖慢速度。固定输入尺寸的静态图是性能最好的形态。注意力模块实现方式。ViT 里最耗时的就是多头注意力。某些 PyTorch 实现用了einops、flash-attn等高级库,导出 ONNX 时代码结构不好,性能直接崩盘。速度异常时用图优化工具看结构是重要的排查手段。有一个很典型的对比案例:同一套 ViT-Base,把头数num_heads从 12 改成 8,序列长度从 196 改成 64,ONNX 导出后的计算图节点数相差三成。这说明模型设计阶段的超参选择对部署性能有直接影响。5.3 用 onnxruntime 自带的 profiling 工具定位瓶颈如果你还是不满足,或者碰到诡异的性能退化,一定要会用 onnxruntime 自带的分析工具。开启 profiling 的方法很简单:so = ort.SessionOptions() so.enable_profiling = True sess = ort.InferenceSession('end2end.onnx', so) # 跑一次推理 sess.run(...) # 结束后会生成一个 .json 文件接着用 Chrome 的chrome://tracing加载这个 JSON 文件,就能看到每个算子消耗的时间和执行顺序。哪一层慢、哪里多了一堆奇怪的Cast节点、哪个Transpose偷走了大量时间,全部一目了然。学会这一手,比盲目调参有效十倍。6. 常见问题与排查方法:我踩过的十个坑,一次说清从 PyTorch 模型到 ONNX Runtime 推理,出错的地方不少。按出现频率排序,把最常见的坑和排查步骤整理如下:6.1 模型导出的常见报错报错现象根因分析解决方案Unsupported operator aten::xxxPyTorch 实现里的某个算子 ONNX 导出器不认识更新 PyTorch 和 mmdeploy 版本;手动导出前把该算子替换为等价算子组合;检查是不是需要更新到更新的 ONNX opsetRuntimeError: Failed to export且错误指向reshapeViT 中动态 shape 导致 reshape 目标维度无法推断设置动态轴(dynamic_axes);用固定输入尺寸;替换内部实现导出成功但推理结果全错数据预处理顺序不对;模型加载前忘记eval模式逐项对比预处理细节,特别是通道顺序和归一化参数;导出前声明model.eval()静态图导出提示None维度数据模型里有未初始化的张量,通常是 batch size 引起的用 dummy input 跑一次 forward 后再导出,而非直接传 NoneGraph must be in single static assignment (SSA) form模型图中存在变量先赋值后重新赋值的张量查看是否有循环内的张量复用;用 trace 而非 script6.2 onnxruntime 推理端常见问题报错现象根因分析解决方案Invalid Symbolic Shape模型含动态轴,session 初始化时无法确定输出形状推理前固定输入 shape;或改用动态 shape 配置;或设置session_options.SetDynamicShape相关选项推理时出现Segmentation fault常见于 C API 中输入张量大小与实际数据不同;内存未对齐;tensor 生命周期提前释放检查CreateTensor时传入的数据指针和大小是否一致;确保input_data在Run前一直存活No implementation found for ...图里包含后端不支持的算子,或后端版本太旧升级 onnxruntime;换用支持该算子的执行提供器(如 TensorRT)或等待算子实现升级Python 接口运行很慢,C API 不慢Python 调用 overhead 在批量小的时候体现明显使用IOBinding减少 Python 层 GIL 开销;增加 batch;把多次推理合并为一个 request图像预处理后输入值范围不对模型期望 0~255 范围的图像输入,你归一化到了 0~1;或反之检查训练时的图像标准化方式,从配置文件的 pipeline 中读取原始值,严格保持一致6.3 一个必须强调的老生常谈:版本一致性很多问题最终都能归结到版本不匹配。onnxruntime 里的算子实现即使算子名相同,不同版本的行为也差异很大。ONNX 的 opset 从 11 到 20,版本越高,支持的算子越丰富,但对老库的兼容性反而可能出问题。如果你不是必须用最新特性,我建议锁死一个稳定组合:PyTorch 1.13 + ONNX 1.13 + onnxruntime 1.14,这组配对我实测下来兼容性最佳。如果手头环境已经换代,至少保证 ONNX 模型的 opset 版本不能高于 onnxruntime 支持的版本。一旦出现算子解析失败,优先检查这个不变量。7. 给 ONNX 模型加个加速 Buff:ONNX Runtime 的线程与内存配置上面提到性能优化时,线程和内存配置值得单独拎出来展开。很多人图省事,直接把默认配置跑完就结束了,结果服务端并发上来之后延迟翻了倍,这时才发现运行时配置有多重要。7.1 SessionOptions 里几个关键开关ONNX Runtime 支持一个专门的配置结构体OrtSessionOptions,其中的参数决定运行时如何调度计算图。SetIntraOpNumThreads:控制单个算子内部的并行线程数。矩阵乘法这类算子内部并行度高,适用于多线程,但线程太多会引入开销。经验值:设置为物理核心数的一半到三分之二。SetInterOpNumThreads:控制计算图中不同算子之间的并行执行。如果你的模型算子是串行依赖关系强,这个参数作用不大;如果模型图结构中有明显的并行分支,把它调大就有收益。SetGraphOptimizationLevel:三档可选,从ORT_DISABLE_ALL到ORT_ENABLE_ALL,直接决定运行时是否对计算图做算子融合。务必设置为ORT_ENABLE_ALL。还有一个容易忽略的是内存分配器。默认情况下 onnxruntime 使用 arena 分配器,适合长生命周期场景。但如果你的推理进程会频繁创建和销毁 session,arena 的内存池反而会造成内存不回收的假象,看起来像泄漏。这时可以设置session_options.AddConfigEntry("session.use_arena", "0")关掉 arena。7.2 遇到批量推理时该如何调整如果你的业务不是一张一张图推理,而是需要一次处理多张,有两种姿势可选。第一种是普通 batch 推理,直接把输入做成(N, 3, H, W)的 N 张图一起送入。这样做吞吐率高,但延迟可能被最慢的那张图拖累。第二种是并发推理,在同一 session 上开多个线程各自跑不同输入。ONNX Runtime 的 session 是线程安全的,多个线程同时调用Run不会冲突。这在 Python 接口上,可以结合concurrent.futures.ThreadPoolExecutor用;在 C API 下,只需各自管理自己的Ort::Value。实测下来,并发推理比单纯增 batch 更灵活,适合服务化场景。8. 工程化落地的几个补充细节8.1 用 onnx 工具库做模型体检模型导出来后,在交给 onnxruntime 之前,建议先做一次全面的模型体检。我在导出完成后用各工具跑了一遍,意义很大。import onnx from onnx import checker, helper, shape_inference model = onnx.load("./work_deploy/end2end.onnx") checker.check_model(model, full_check=True) # 做一次形状推断,能暴露隐藏的维度错误 inferred_model = shape_inference.infer_shapes(model) onnx.save(inferred_model, "./work_deploy/end2end_inferred.onnx") # 统计节点类型,看是否有意外出现的算子 from collections import Counter node_types = Counter(n.op_type for n in model.graph.node) print(node_types.most_common(20))最后这行统计很有用。我遇到过一种情况,模型里出现了大量的Cast节点,原因是 ViT 的某些算子对精度有要求,被强制转换成了double。这种精度提升对分类任务没有明显收益,却让 CPU 推理时间翻倍。发现之后,我修改了导出配置,统一用 FP32,速度立刻就回来了。8.2 模型输入的 meta 信息管理服务化部署时,把模型文件和一些元信息放在一起管理很重要。模型除了本身的 ONNX 图,还有输入输出节点名称、图像尺寸、归一化参数等关联信息。mmdeploy 的--dump-info参数会生成pipeline.json,里面包含了上述信息,但格式不直观。我的做法是自己在模型旁边单独放一个 JSON 描述文件:{ "model": "end2end.onnx", "input": { "name": "input", "shape": [1, 3, 224, 224], "dtype": "float32" }, "output": { "name": "output", "shape": [1, 1000], "dtype": "float32" }, "preprocess": { "mean": [123.675, 116.28, 103.53], "std": [58.395, 57.12, 57.375], "color_order": "BGR", "resize": "224x224" } }有了这个文件,模型文件本身就可以做到平台无关。换任何语言、任何框架来接 ONNX 模型时,参数的读取都有了统一来源,不会出现模型文件换台机器就跑不起来的窘境。8.3 单元测试与精度比对方案模型的精度验证是部署里最容易马虎又最不该马虎的环节。我的精度验证标准方案是:选 100 张测试图,分别用 PyTorch 原模型和 ONNX Runtime 推理。对每张图比较两边的 top-1 类别是否一致。对置信度分数,计算平均绝对误差(MAE)。误差通常在1e-4到1e-5量级,这是 FP32 精度在不同算子实现下的正常波动。如果误差超过1e-2,基本可以断定是预处理、数据格式或算子精度出了问题。这个小脚本虽然简单,但有效。我记得有一次导出的模型在测试集上准确率掉了 3 个百分点,排查了整整半天,最后发现图像归一化的std写成了训练时使用的mean,类似这种错误在视图比对时能很快暴露问题。9. 再延伸一步:为 ONNX 推理加个“加速 Buff”,动态输入与量化探索如果前面都没问题,还有两个方向可以继续挖掘,把 ViT 模型的部署性能再往前推一点。方向一:动态输入尺寸的支持。ViT 在导出时如果用动态轴配置,就可以输入任意尺寸的图片,而不必强制 reshape 到 224x224。这对真实场景很有意义,客户发来的图片构不成正方形,直接捏成 224x224 会损失大量信息。动态轴的配置在 mmdeploy 的配置文件中可以指定:dynamic_axes = { "input": {0: "batch", 2: "height", 3: "width"}, "output": {0: "batch"} }配置之后重新导出,就能支持任意尺寸输入。代价是推理时会有额外的 shape 推断开销,每次输入尺寸如果和上次不同,ONNX Runtime 需要重新编译部分算子,延迟会略有上升。如果你场景里图片尺寸变化不频繁,动态输入带来的精度收益远大于这点开销。方向二:量化。CPU 部署模型时,INT8 量化是提升速度最有效的手段。ONNX Runtime 提供onnxruntime.quantization工具包,静态量化需要一批校准数据来计算激活值的动态范围:from onnxruntime.quantization import quantize_static, QuantType, CalibrationDataReader class DataReader(CalibrationDataReader): def __init__(self, data_paths): # 实现 get_next 方法,返回一批校准样本 pass quantize_static( model_input="end2end.onnx", model_output="end2end_int8.onnx", calibration_data_reader=DataReader([...]), quant_format=QuantType.QInt8 )量化后的 ViT 模型在 CPU 上通常能获得 2~3 倍的速度提升,精度损失在分类任务中一般在 1%~2% 以内。这条路线如果业务需求有压测要求,非常值得深入研究。我的总结体会这条链路跑通之后,往后遇到任何 OpenMMLab 系模型的部署需求,基本都能沿着一套模式完成:模型准备、导出 ONNX、选择推理后端、处理输入输出、性能调优。整个过程看下来,真正费时间的从来不是写代码,而是版本匹配、预处理一致性、动态 shape 这些问题。一个朴素的想法是,跨框架部署的本质不是“模型格式转换”,而是把训练时的所有假设——包括数据的进出、算子的边界、精度容差——原原本本地翻译到另一端。只要抓住了这一条,很多诡异问题其实都能一眼看穿。最后再分享一个我个人的小习惯:每导出一个新模型,我都会把环境依赖的版本号写进一个 requirements 文件里,和模型文件一起保存。三个月后再回来看,依赖已经大变样了,但只要有那份文件,任何方法都能快速复现当时的部署环境。这个习惯帮我避过太多次的版本灾难,建议你从现在开始也这样做。