企业官网建设流程全解析

1. OpenMMLab里搞到ViT：别再傻傻手写Transformer

做视觉模型部署的人应该都有体会：训练容易，部署难。PyTorch里跑ViT（Vision Transformer）很流畅，但到了生产环境、CPU服务器或者边缘设备上，eager模式的调度开销和算子碎片化问题就会被无限放大。我最近在一个图像分类项目里就把OpenMMLab体系训练的VisionTransformer模型导出成ONNX，再用onnxruntime做推理，整条链路跑通之后，推理速度提升非常明显。这篇文章就顺着这个主题，把从OpenMMLab拿模型到ONNX落地推理的完整过程梳理一遍，适合正在做ViT部署、或者刚接触模型导出的同学参考。

先说明一下OpenMMLab和ViT的关系。OpenMMLab不是一个单一模型仓库，而是一套算法体系。ViT在里面的核心归属是MMPretrain（它继承自原来的MMClassification），同时MMDetection、MMSegmentation等下游库里也会用ViT当backbone。所以第一步就不是写模型结构，而是从这套体系里把带权重的ViT正确取出来。

1.1 在MMPretrain里找backbone的正确入口

MMPretrain的模型库分类很规整，按backbone名去索引就行。我习惯直接用它的命令行工具查看可用模型：

mmpretrain list models | grep vit

这条命令会列出所有ViT相关配置，包括预训练权重和训练设置。列出来的模型里，vit-base-p32、vit-base-p16、vit-large-p16这些是主力。要注意命名后缀，比如_ft_64e_in21k_1k表示先在ImageNet-21k上预训练、再在1k上微调64轮，这类权重做迁移和部署都更稳。

拿到模型最简单的方式是直接用get_model：

from mmpretrain import get_model model = get_model('vit-base-p16_ft-64e_in21k-1k', pretrained=True) model.eval()

pretrained=True会从线上权重库拉预训练模型，第一次运行会下载文件。如果你处在离线环境，可以手动下载后放在指定缓存目录。这个接口返回的不只是一个backbone，而是带着分类头的完整分类器，包含backbone、neck和head三部分。

还有一个常见需求：只导出backbone做特征提取。这时候get_model就不太合适，因为分类头会干扰部署，你并不需要那层MLP分类器。更干净的做法是从MMPretrain的backbones模块里直接构造backbone：

from mmpretrain.models.backbones import VisionTransformer import torch backbone = VisionTransformer( arch='b', img_size=224, patch_size=16, frozen_stages=12, norm_eval=True, out_type='cls_token', ) backbone.init_weights()

这里解释几个参数的用意。arch='b'是base规模定义，隐藏维度768、12层Transformer、12个注意力头，和ViT-B/16对应。out_type='cls_token'让backbone输出CLS token，这个token本身已经聚合了全局特征，适合直接接分类头用。norm_eval=True表示推理时LayerNorm的参数不参与BN统计更新，这也是部署时必须锁死的状态。

1.2 把训练模型切成纯推理模型

很多人在导出时出问题，根本原因不是导出代码写错了，而是模型里混进了训练相关的组件。MMPretrain的分类器默认带neck和head，head里有Dropout、全连接层，甚至还有loss模块。这些在推理阶段完全用不到。

我的做法是拿到模型后重构一个前向接口，只保留到backbone的输出，再包一层ClassiferHead也行，但更省事的是直接切掉head：

from mmpretrain import get_model full_model = get_model('vit-base-p16_ft-64e_in21k-1k', pretrained=True) backbone = full_model.backbone

注意，这样切出来的backbone输出是一个tuple，里面可能是(cls_token, patch_tokens)或者只有cls_token。导出ONNX时torch.onnx.export对tuple输出会有额外的处理开销，建议再包一层，让它只返回单一tensor：

class ViTBackboneExport(torch.nn.Module): def __init__(self, backbone): super().__init__() self.backbone = backbone def forward(self, x): out = self.backbone(x) if isinstance(out, (tuple, list)): return out[0] return out export_model = ViTBackboneExport(backbone).eval()

这样导出图就干净很多。后面做onnxruntime推理时，输出就直接是CLS特征或者logits，省掉很多纠缠。

2. ViT结构里的算子，在导出时到底映射成了什么

ViT在PyTorch里写起来很清爽，几行attention代码就搞定了。但到了ONNX图里，它会被展开成大量细粒度算子，这也是部署后图体积大、优化难度高的原因。理解这一层，才能明白后面做图优化的必要性。

2.1 ViT主体结构回顾

一个标准ViT-B/16可以被拆成五个逻辑模块：

• PatchEmbed：把224×224×3的图像切成14×14=196个16×16的patch，每个patch通过一个Conv2d(3, 768, kernel_size=16, stride=16)映射成768维向量。
• Positional Embedding：把196个patch token加上一个可学习的position embedding，再拼接一个CLS token，最终输入序列长度是197。
• Transformer Encoder块×12：每个块内部是LayerNorm、Multi-Head Attention（QKV映射、注意力分数计算、加权求和）、残差连接、LayerNorm、MLP（两个全连接层，中间有GELU激活），再接残差。
• 最后的LayerNorm：给输出CLS token做归一化。
• 分类头（部署时可去掉）：一个Linear层，把768维映射到类别数。

导出时，上面的每一步都会对应到ONNX的某个算子。PatchEmbed对应Conv算子，position embedding对应Add算子，LayerNorm对应LayerNormalization（或者更低opset下的ReduceMean、Sub、Pow、Div组合），GELU在ONNX里没有原生算子，torch导出时一般用Erf配合加减乘除来实现，Multi-Head Attention里的view、permute会展开成Reshape、Transpose、Squeeze、Unsqueeze这些算子。

2.2 每个子模块导出的ONNX算子形态

我导出过一次ViT-B/16后专门用Netron看过计算图，印象很深。一个MatMul加Mul再加Softmax这种在PyTorch里一行搞定的注意力计算，在ONNX里会变成Reshape → Transpose → MatMul → Mul → Softmax → MatMul → Transpose → Reshape，整整一大串。

各模块对应的算子大致如下表：

这张表很有用，你导出后用Netron拉图对照，基本能一眼看出哪个算子来自哪个模块，排查速度会快很多。

2.3 为什么导出后算子数会膨胀

PyTorch模型和ONNX模型在算子粒度上不对齐。PyTorch里permute只改一下tensor的元数据，不真正搬运数据，但在ONNX里Transpose是实打实的数据重排算子。同理，view在PyTorch里也几乎是零开销，但在ONNX里会生成一个Reshape节点。

ViT的attention部分集中了大量的view/permute，一次多头拆分会生成两三个Reshape + Transpose，12层Transformer叠下来，算子数轻松翻一倍以上。我导出的原始ViT-B/16，ONNX文件里Reshape+Transpose节点占了将近40%，这些节点不搬数据、不计算，纯粹是形状转换，但推理引擎要一个一个执行，严重拖慢速度。

所以导出后的图优化不是锦上添花，而是必须做的一步，下一节就专门讲。

3. torch.onnx.export导出实操：参数、动态维度、避坑

说完了模型准备和结构映射，进入正题：跑导出代码。

3.1 最简导出代码

假设你已经有了一个干净的export_model（只输出单一tensor），用下面这段代码就能导出：

import torch dummy_input = torch.randn(1, 3, 224, 224, device='cuda') export_model = export_model.cuda().eval() with torch.no_grad(): torch.onnx.export( export_model, dummy_input, "vit_b16.onnx", input_names=["input"], output_names=["output"], dynamic_axes={ "input": {0: "batch"}, "output": {0: "batch"}, }, opset_version=17, do_constant_folding=True, )

几个细节特别说明一下。

export_model的eval()必须显式调用，最好再包一层torch.no_grad()。ViT里虽然没有BN，但LayerNorm在训练和推理模式下行为一致，就怕模型里混着dropout，训练模式下dropout会随机丢token，导出的图数值不稳定。

dummy_input的shape必须和模型实际输入一致。我这里用1, 3, 224, 224，如果你的模型用的是OpenMMLab的RandomResizedCrop逻辑，输入可能是3, 224, 224，注意别带反了。

do_constant_folding=True这一步很关键，它会把模型里的常量计算（比如position embedding的Add）在导出时就折叠成常量。如果设成False，ONNX图里会留下一堆纯常量计算的算子，推理时每次都要重复计算一遍。

3.2 dynamic_axes不是随便开的

很多教程告诉你要把输入输出的维度都设置成动态，但ViT里有个特殊约束：position embedding和patch数量是强绑定的。

ViT的输入序列长度由img_size / patch_size决定，224/16=14，序列长度196。Position embedding是和这个长度一起训练的，如果你把height、width设为动态维度，模型推理时输入一个384×384的图，patch数就变成24×24=576，超出了position embedding的形状，直接报错。除非你用了支持任意分辨率的交叉注意力或者对position embedding做了插值，否则动态H/W在ViT导出里没有意义。

所以我的建议是：动态维度只开batch，H和W保持静态。这样既满足了批量推理需求，又避免了序列长度对不上的问题。

3.3 opset版本怎么选

opset就是ONNX的算子集版本，不同版本支持的算子类型和精度不一样，对ViT影响最大的有两个地方：LayerNorm和GELU。

在opset 13及以下，LayerNorm没有原生算子，torch会把nn.LayerNorm拆成ReduceMean→Sub→Pow→Div→Mul→Add一长串，节点多、中间结果多，推理时数值上还会有点小尾差。

opset 17开始，LayerNormalization算子已经成熟，PyTorch导出器也会优先走原生算子，导出图一下子干净很多。

opset 20开始出现了原生Gelu算子，但onnxruntime对opset 20的支持还在迭代中，部分老版本runtime会遇到算子兼容问题。

我实测下来，opset 17是当前最稳的选择，ONNX Runtime完全支持，LayerNorm又是原生算子，导出精度和推理速度都比较理想。如果你一定要和某个只支持低opset的推理框架对接，那再降到opset 13也不迟，但要接受LayerNorm被拆开的代价。

4. 导出后的图优化与形状问题处理

拿到第一次导出的vit_b16.onnx，千万别直接拿去跑推理。先看一下文件大小和算子数，你就明白为什么要处理了。

4.1 onnxsim一键清理冗余算子

onnx-simplifier（pip包名onnxsim）是目前最主流的ONNX图优化工具，专门处理这种"PyTorch风格算子碎片化"问题。它的原理是：输入ONNX图，逐个节点做常量折叠、形状推导、死节点消除，能合并的Reshape和Transpose合并，能静态化的动态维度静态化。

安装和用法：

pip install onnxsim python -m onnxsim vit_b16.onnx vit_b16_sim.onnx \ --overwrite-input-shape input:1,3,224,224

--overwrite-input-shape是为了让simplifier在固定shape下做符号形状推导，这一步对ViT尤其重要。因为onnxsim默认会保留dynamic_axes里设置的动态维度，但ViT的图经过position embedding和reshape后，很多节点之间的形状其实是静态的，固定输入shape后才能把这些约束传导下去，合并掉更多冗余节点。

跑完之后对比一下：

这个数据是我在ViT-B/16上的实际结果，Reshape和Transpose减少了一大半，推理引擎不用再空转那些形状转换节点，速度直接上一个台阶。

4.2 静态shape还是动态shape的取舍

导出的onnxsim版本默认绑定batch=1的静态shape。如果你的部署场景就是单张图推理，静态shape是性能最优解，因为引擎可以做更多的算子融合优化。

但如果你需要batch推理，建议在导出时就用dynamic_axes打开batch维度，再通过onnxsim生成一个batch可变的模型。注意，onnxsim处理动态shape时有个特点：如果input:1,3,224,224里指定了固定batch=1，它就会把batch固定死。所以我一般分两套模型：

• 线上单图服务：静态shape模型，CPU上单张推理延迟最低。
• 批量离线处理：batch动态模型，GPU上充分吃满显存并行。

这两个场景别混用，混用反而两边都优化不好。

4.3 LayerNorm算子在低opset下的数值小尾巴

如果你选择的opset低于14，还要特别做一次精度验证。LayerNorm被拆开后，ReduceMean求均值、Pow求方差，中间结果的浮点运算顺序和原始PyTorch不完全一致，小数点后第6-7位会出现差异，放大到深层网络后可能影响最终logits。

我的判断标准是尽可能用原生LayerNormalization算子，让它和PyTorch内部实现在算子层面保持一致。如果推理环境强制要求低opset，那我就把LayerNorm的数值误差当作已知偏差，在精度验收时留好余量，别卡死阈值。

5. onnxruntime推理：会话配置与输入预处理对齐

模型文件准备好了，接下来就是在onnxruntime里把它跑起来。这一节最容易出问题的地方不在ONNX模型本身，而在推理环境配置和输入预处理。

5.1 环境准备与Session配置

先装依赖：

pip install onnxruntime # 或者 GPU 版本 pip install onnxruntime-gpu

判断有没有装对，直接看provider列表：

import onnxruntime as ort print(ort.get_available_providers())

如果只有CPUExecutionProvider，说明装的是CPU版本；GPU版本会列出CUDAExecutionProvider甚至TensorrtExecutionProvider。

推理Session的配置直接影响性能，我最常用的配置如下：

import onnxruntime as ort import numpy as np sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.intra_op_num_threads = 8 sess_options.inter_op_num_threads = 1 sess = ort.InferenceSession( "vit_b16_sim.onnx", sess_options, providers=["CPUExecutionProvider"], ) # 查输入输出信息 for inp in sess.get_inputs(): print(f"input: {inp.name}, shape={inp.shape}, type={inp.type}") for out in sess.get_outputs(): print(f"output: {out.name}, shape={out.shape}, type={out.type}")

graph_optimization_level=ORT_ENABLE_ALL是默认值其实也是最高优化级别，它会启用图融合、算子替换、层融合等一系列优化。对于ViT这种算子密集型模型，ORT会自动把连续的Reshape+Transpose+MatMul做融合优化，收益明显。

intra_op_num_threads和inter_op_num_threads这两个参数要分开理解：前者控制单个算子内部的线程数，后者控制图节点间的并行任务数。ViT计算图是一条很深的串行链，Transformer层和层之间有数据依赖，不适合开太多inter_op并行。实践中我把inter_op设为1、intra_op设为物理核心数，效果最好。

5.2 预处理必须对齐OpenMMLab的pipeline

这一步是新手最容易踩的坑。ONNX模型不会记住你的预处理参数，训练时的normalize均值、标准差、resize方式、通道顺序都必须在推理代码里手动复现。OpenMMLab的标准ImageNet预处理pipeline是：

1. 读图
2. Resize到256×256
3. CenterCrop到224×224
4. 通道变成RGB，归一化：像素值/255，然后减mean除std，mean=[0.485, 0.456, 0.406]，std=[0.229, 0.224, 0.225]
5. HWC转CHW、转float32

对应到推理代码：

import cv2 import numpy as np def preprocess(image_path): img = cv2.imread(image_path) # BGR, HWC img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_LINEAR) h, w = img.shape[:2] start_h = (h - 224) // 2 start_w = (w - 224) // 2 img = img[start_h:start_h + 224, start_w:start_w + 224] img = img.astype(np.float32) / 255.0 mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.float32) img = (img - mean) / std img = img.transpose(2, 0, 1) # HWC -> CHW img = np.expand_dims(img, axis=0) # 加batch维 return np.ascontiguousarray(img)

这里有个细节要强调：CenterCrop的起点必须是(resize后尺寸 - 目标尺寸) // 2。OpenMMLab用它自己的CenterCrop实现，PyTorch的transforms也一样，但如果你直接用torchvision.transforms.CenterCrop去处理numpy，边界条件可能不同。我见过有人用img_center做像素中心裁剪，写起来差不多，但角点位置差一两个像素，导致精度下降，排查了很久。

还有resize的插值方式。OpenMMLab默认用cv2.INTER_LINEAR（双线性），有些部署代码图省事用INTER_NEAREST（最近邻），两者差异在224×224输入下大约会带来千分位的精度损失。

预处理完，直接喂给session跑：

input_data = preprocess("test.jpg") output = sess.run(None, {"input": input_data})[0] print(output.shape, output[:, :5]) # 打印top-5逻辑回归值

如果模型输出的是分类logits，一般还要做一个softmax再取top-1。OpenMMLab分类模型head里通常已经带softmax了，但如果你在导出时是切出来的backbone，就需要自己在推理侧补softmax，别漏。

5.3 CPU线程参数这样调

CPU推理线程调优，我的经验是分两步。

第一步看机器物理核数。ViT是大模型，算子内部并行才有效率。比如一台8核16线程的CPU服务器，我设intra_op_num_threads=8而不是16，因为超线程对计算密集的MatMul帮助不大，反而会增加线程切换开销。

第二步实测对比。把上面的配置改几次，用同一样本跑100次推理，取平均时延。我实测的一个结果是：

最后一行很有意思，inter开太高之后图级并行任务调度开销超过了收益。所以默认建议抓住两个关键词：intra_op多、inter_op少。

6. 精度验证与性能对比：这一环不能省

模型导出完、推理跑通了，很多人的项目就结束了。但部署上线前，还有最关键的验证工作：确保ONNX输出和PyTorch原模型一致，以及确认性能提升是实打实的。

6.1 数值一致性怎么量化

数值验证方式很简单：同一张图，同一个预处理结果，分别喂给PyTorch模型和ONNX模型，对比输出张量。

import torch import numpy as np torch_input = torch.from_numpy(input_data).cuda() with torch.no_grad(): torch_output = export_model(torch_input).cpu().numpy() ort_output = sess.run(None, {"input": input_data})[0] # 余弦相似度 def cosine_sim(a, b): a = a.flatten() b = b.flatten() return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12) # 最大绝对误差 max_abs_diff = np.abs(torch_output - ort_output).max() print(f"cosine_sim: {cosine_sim(torch_output, ort_output):.8f}") print(f"max_abs_diff: {max_abs_diff:.2e}")

我通常会跑三张代表性图片：一张正常光照、一张偏暗、一张带噪点。ViT的数值稳定性整体不错，f32导出下余弦相似度一般在0.9999以上。如果这个值低于0.999，先检查预处理，再检查模型是不是意外切了dropout，最后再考虑算子精度问题。

还有一种情况：导出时开了dynamic_axes的batch维度，但PyTorch侧输入是batch=1。建议batch维度也测一下，比如输入batch=4做一次对比，防止动态batch引入shape异常。

6.2 性能测试方法与实测结果解读

性能测试要做预热。第一次推理会包含显存/内存分配、模型加载、图的JIT编译等开销，不能计入统计。正确做法是先跑几次预热，再计时。

import time # 预热 for _ in range(10): sess.run(None, {"input": input_data}) # 正式计时 times = [] for _ in range(100): t0 = time.perf_counter() sess.run(None, {"input": input_data}) times.append(time.perf_counter() - t0) times.sort() print(f"avg: {np.mean(times)*1000:.2f} ms") print(f"p50: {times[50]*1000:.2f} ms") print(f"p95: {times[95]*1000:.2f} ms")

PyTorch侧做对比时也要预热，并且要确认模型在eval模式、torch.no_grad()下，否则反向图计算会带来额外开销。

我实测Vit-B/16在相同CPU上的数据大概是这样的：

onnxsim前后差了32%，这个数字再次印证了VIT类模型图优化的必要性。GPU环境下的对比更夸张，PyTorch eager在V100上大约15ms，ORT用CUDA provider能到10ms以内，但前提是provider配好。

最后提醒一句：性能测试结果和硬件高度相关，上面的数据只能当参考量级，不同CPU上差异很大。最重要的不是绝对数字，而是对比出ONNX相对PyTorch的提升，以及onnxsim带来的额外收益。

我个人做这类型部署踩过不少坑，最大的体会是：ONNX导出这件事，模型结构知识比重远大于onnx库知识。把ViT里每个PyTorch模块映射到ONNX算子吃透了，后面所有问题都能顺藤摸瓜找到根因。特别是Reshape和Transpose的膨胀问题，几乎所有Transformer类模型导出都会遇到，用onnxsim处理是一劳永逸的解法。下一步我想把这条链路继续延伸到量化部署和TensorRT后端，等实测数据出来了再来分享。