最近有个朋友问我他手里那套基于 YOLOv5 的检测系统在服务器上跑得好好的现在想推上移动端但换了轻量版之后精度又不太满意。问我有没有什么方案能在保持速度基本不变的前提下把 mAP 再往上拉一截。我给他推荐了我之前折腾过的一条路子把 MobileViTv1 的核心模块融进 YOLOv5 的 backbone 里训练出一个混合型检测模型。实测下来在几乎不牺牲推理速度的情况下mAP 能提升 1.8 到 2.5 个点左右整个模型的体积增加也不多移动端部署完全能接受。这个项目对我来说挺有代表性的它不是那种“换个大模型提精度”的思路而是在轻量化的框架里做精细化改造。整个过程涉及模型结构改造、训练策略调整、ONNX 导出、移动端算子适配这些环节坑踩了不少收获也很大。这篇就完完整整地把我的设计思路、代码改动、实验数据、踩坑记录都写下来。想在自己项目里复现的朋友可以直接按着步骤走就算你不是做检测方向的看一个 CNN 和 Transformer 怎么在目标检测器里低成本融合也能有不少启发。1. 融合思路为什么是 YOLOv5 MobileViTv11.1 YOLOv5 的 backbone 真的到瓶颈了吗先说个可能不太好听的事实很多人在移动端部署 YOLOv5用的还是最原始的 YOLOv5s 结构backbone 里的 C3 模块从头到尾全是卷积操作。C3 的结构本身没什么问题CSP 的拆分思想也足够高效但它终究是局部算子。卷积核再大也就是 3x3、5x5 这种窗口感受野始终有限。到了网络高层特征图分辨率变小了一个格子对应的原图区域变大了可是每个格子依然只能看到自己周围的局部上下文没法把整张图范围内的语义关联起来。我打个比方一张照片里有两个长得差不多的行人一个在左下角一个在右上角中间隔着很远的背景。CNN 的 backbone 在处理这两个人的特征时是完全独立的两个分支它们之间没有信息交换。如果模型能提前知道“这张图里有两个外形相似的物体”它对其中一个的判断就会更有把握。Transformer 的自注意力机制恰恰能解决这个问题它让每个位置都能“看到”全图其他位置的信息。MobileViTv1 就是在这种思路上诞生的。1.2 MobileViTv1 到底赢在哪MobileViTv1 是轻量级视觉 Transformer 里的一个代表作。它和原版 ViT 那种“上来就把整张图切成大 patch然后一路全局 Transformer”的做法完全不同MobileViT 只在局部 patch 之间做自注意力patch 与 patch 之间的信息交换仍然依靠卷积来完成这样就把计算复杂度从 O(n²) 降到了接近纯 CNN 的水平。具体理解起来也不难。你可以把一张特征图当成一个棋盘MobileViT 里面的 Transformer 编码器只在每个小方格内部做注意力让方格内部的像素之间建立长程依赖而方格和方格之间通过卷积的滑动窗口来沟通。这样每一个特征点既拥有了全局视野又不会像原版 ViT 那样带来恐怖的算力消耗。更重要的是MobileViT 的输入通道设计沿用了 MobileNet 系列的倒残差结构计算量本来就很低。把这种模块嵌进 YOLOv5理论上能做到“速度不明显变慢但精度提升”。我在这个项目中用的版本是 MobileViTv1而且只引入它最核心的 MobileViTBlock不把整个 MobileViT 骨干搬过来。原因下面会详细说。1.3 两种融合策略到底怎么选当时摆在我面前有两条路线。第一条是“局部替换”保留 YOLOv5 原有的骨干骨架把 backbone 里的部分 C3 模块整体替换成融合了 MobileViTBlock 的新模块。第二条是“整体换血”直接把 YOLOv5 的整个 backbone 全部换掉让 MobileViT 的 Stem 部分接管输入然后再接 YOLOv5 的 neck 和检测头。从效果上限来看整体换血理论上更接近 MobileViT 论文里的完整结构但工程上麻烦很多。首先你要重新对齐下采样的 strideYOLOv5 的 backbone 要输出 8 倍、16 倍、32 倍三个尺度的特征给 neckMobileViT 原版结构的输出不见得刚好匹配其次原来的 Focus 或者 Stem 部分的预训练权重全部作废得从 ImageNet 预训练权重里重新映射工作量很大。所以我最终选了第一条路线保留 YOLOv5 主结构只在高层语义阶段引入 MobileViTBlock让 Transformer 在高层专门负责全局建模。具体到我的版本我只替换了 backbone 中 P3、P4 两个阶段的部分 C3P2 和 P5 阶段保持原版 C3 不变。这么做还有个好处可以单独做消融实验对比“有 MobileViT 和没有 MobileViT”各自带来了多少收益不会一上来就是黑盒。2. 核心模块拆解与代码改造2.1 MobileViTBlock 的五个关键步骤如果你想改代码得先把这个模块的结构看透。MobileViTBlock 的 forward 整体可以拆成几大步。第一步先用一个 1x1 卷积做通道变换把通道数升到一个中间维度为后面的局部建模和全局建模准备一个更丰富的特征空间。第二步用一个 3x3 的深度可分离卷积提取局部信息这一步和 MobileNet 里的 depthwise conv 完全一样计算量非常小。第三步是核心把当前特征图按 patch_size 切成一个个小格子每个格子里所有像素被展开成一个 token 序列然后在这个序列上做标准的 Transformer 编码包括 LayerNorm、多头自注意力、FFN同时加一个可学习的位置编码。第四步把处理完的序列再还原成特征图的形状这个操作对应着 PyTorch 里的 fold和第三步的 unfold 是配套的。最后第五步再用 1x1 卷积把通道数降回来并且和输入做残差相加。这里有个特别需要注意的点unfold 和 fold 要求特征图的宽高必须能被 patch_size 整除否则还原后尺寸会对不上。我在实验里一般用 patch_size2并提前把输入图像 pad 成偶数尺寸这样不会出问题。2.2 在 YOLOv5 的 common.py 里新增类YOLOv5 的模型定义比较工程化所有基础模块都在models/common.py里比如 Conv、Bottleneck、C3、SPPF。要在 YOLOv5 里用 MobileViTBlock第一个要做的事就是在 common.py 里加类和注册逻辑。下面是我加进去的核心简化版本感兴趣的可以参考class MobileViTBlock(nn.Module): def __init__(self, dim, depth, channel, kernel_size3, patch_size2): super().__init__() self.patch_size patch_size self.dim dim # 局部特征提取 self.local_rep nn.Sequential( Conv(dim, dim, 1), Conv(dim, dim, kernel_size, groupsdim), ) # Transformer 编码层 self.transformer nn.ModuleList([ TransformerEncoderLayer(dim, ...) for _ in range(depth) ]) self.pos_embed nn.Parameter(torch.zeros(1, dim, patch_size * patch_size)) def forward(self, x): y self.local_rep(x) B, C, H, W y.shape # resize 到 patch 序列 y y.reshape(B, C, H // self.patch_size, self.patch_size, W // self.patch_size, self.patch_size) y y.permute(0, 1, 2, 4, 3, 5).reshape(B, C, -1, self.patch_size * self.patch_size) # 在每个 patch 内做 Transformer 编码 # fold 回原尺寸 # 与输入做残差并输出 return out x不过我要提醒一下上面只展示了核心思路真正的完整实现还要处理 unfold 时的 padding、LayerNorm 的顺序、多头注意力的 mask、残差路径的通道对齐等细节。另外我这里的 MobileViTBlock 会和 YOLOv5 的 C3 结构组合成一个新模块我把它命名为C3MobileViT内部逻辑是输入先进一个 1x1 卷积然后走若干次 MobileViTBlock最后再和 shortcut 分支相加。这个思路和原版 C3 一致只是把中间的 Bottleneck 换成了 MobileViTBlock。2.3 yaml 模型配置文件怎么改YOLOv5 的模型结构全部由 yaml 配置驱动改完 common.py 之后要新建一个yolov5s_mobilevit.yaml。我当时的配置很克制只在骨干的第三、第四阶段做替换backbone: - [-1, 1, Focus, [64, 3]] - [-1, 1, Conv, [128, 3, 2]] - [-1, 3, C3MobileViT, [128, 3, 2]] - [-1, 1, Conv, [256, 3, 2]] - [-1, 6, C3MobileViT, [256, 3, 2]] - [-1, 1, Conv, [512, 3, 2]] - [-1, 9, C3, [512]] - [-1, 1, Conv, [1024, 3, 2]] - [-1, 3, C3, [1024]] - [-1, 1, SPPF, [1024, 5]]这里 repeats 我没有完全按照 YOLOv5s 的默认值走而是把 P3、P4 阶段的 C3MobileViT 的层数调大了一些。你要注意parse_model在解析 yaml 时会根据模块名去 common.py 里找对应的 class如果找不到就会报错所以名字一定要对齐。我当时在 yaml 里写的是C3MobileViT没注册成C3_MobileViT排查了半天才发现是大小写下划线的问题。2.4 哪些部分完全不用动很多人一上来就激进觉得 MobileViT 是 Transformer是不是 neck 里的 PANet 也得改成 Transformer我建议不要动。PANet 的 FPN 结构已经很成熟它的主要任务是做多尺度特征融合不需要全局建模检测头更是和 anchor 强绑定的改了就要重配对。所以我的改动范围就是 backbone 的两个 stageneck 和 head 保持原样这直接决定了项目能快速收敛、快速验证。3. 训练配置、消融实验与数据变化3.1 数据集与实验基线的准备实验用的是我手头一个场景数据集里面既有常见物体也有大量小目标大概 4.2 万张按 8:1:1 分成训练、验证、测试。类别数是 18不是 COCO 的 80 类所以我把默认的 nc 和 anchor 都重新算了。锚框计算我直接用了 YOLOv5 自带的聚类脚本针对训练集重新聚类这样比较公平。在训练任何融合模型之前先用原版 YOLOv5s 训了一条 baseline作为后续所有对比的参考。这个步骤一定要做否则你根本分不清 mAP 的提升到底来自 MobileViT还是来自你改动的训练超参。baseline 和融合模型的所有训练超参保持一致这是实验的基本原则。3.2 消融实验从单阶段替换到全面替换我做了三组核心对比。第一组是原始 YOLOv5s第二组只替换 P3 阶段的 C3 为 C3MobileViT第三组替换 P3P4 两个阶段。结果非常典型模型配置参数量GFLOPsmAP50mAP50:95GPU 推理耗时(ms)YOLOv5s baseline7.23M16.50.8520.5186.8只替换 P38.14M18.20.8610.5297.4替换 P3P49.05M20.10.8690.5378.1可以看到只替换 P3 阶段mAP50:95 提升了 1.1 个点参数量增加不到 1M替换 P3P4 之后又提升了 0.8 个点左右。整体来看把 C3 换成 C3MobileViT计算量增加的幅度较小但精度收益是实打实的。我还试过把 P5 阶段也替换成 C3MobileViT结果精度反而降了 0.3 个点。原因可能是 P5 特征图分辨率太低已经接近 anchor 的粒度再做全局注意力容易过拟合。3.3 训练细节别让模型从零开始这个项目里最影响最终效果的是一个习惯我没有让融合模型从随机初始化开始训练而是从原版 YOLOv5s 的骨干权重做热启动。你看 yaml 中 P3、P4 阶段替换成了 C3MobileViT原版 YOLOv5s 在这两个阶段里没有对应模型那些新增 MobileViTBlock 的参数我会额外加载 MobileViT 自己预训练好的权重。方法是对齐 state dict 里面的参数名把能匹配上的复制过来匹配不上的保持默认初始化。实测下来这种初始化方式比完全随机初始化在早期收敛速度快很多最终 mAP 也高 1 个点以上。训练超参我这样设的输入尺寸 640batch size 32单卡epoch 120初始学习率 0.0025用了 3 个 epoch 的 warmupLR 策略用 cosine annealing。开了 EMA但关闭了自动锚框优化因为 anchor 已经在训练前单独对齐过。对比实验里我设置完全一致确保差异只来自模型结构。3.4 没有带来收益的三个尝试这里分享一些反向经验比正向经验的含金量更高。我第一次试的时候把 MobileViTBlock 放到了所有 C3 位置包括浅层 P2。浅层特征图尺寸大比如 160x160 或 80x80全局注意力的序列长度直接爆炸显存一下就不够用了而且精度没有提升。我后来把浅层的序列长度控制住了但训练速度依然慢很多收益却不明显。第二个尝试是把 patch_size 从 2 改成 4想进一步压缩计算量。速度确实快了一点但 mAP50:95 掉了 1.1 个点小目标漏检明显增多。因为 patch 大了每个 patch 内部的点被强行归成一组做注意力局部细节被过度平均了。第三个尝试是不加位置编码。原版 MobileViT 在 patch 内部加了可学习的位置编码我觉得效果不明显就删了结果训练怎么也不收敛loss 一直降不下去。补回去之后立刻恢复正常。这个教训是Transformer 没有位置信息就分不清先后顺序哪怕是在局部 patch 里位置编码也不能省。4. 部署转换与移动端落地4.1 ONNX 导出时最容易翻车的点训练完模型下一步要部署到移动端。不管是用 NCNN、TNN 还是 TensorRT第一步基本都是先把 PyTorch 模型导出成 ONNX。但这个项目里含有一个特殊操作unfold 和 fold。这两个操作在导出 ONNX 时经常被转换成动态 Reshape而 Reshape 又依赖具体输入尺寸一旦你导出的动态维度处理不好后续在推理引擎里就会各种报错。我的做法是导出时直接固定输入尺寸为 640x640不要开动态 H/Wopset 用 12 以上因为低版本的 opset 对 unfold 和 fold 支持很有限。在导出脚本里加上--simplify用 onnx-simplifier 把冗余节点清掉TFLite 和 NCNN 转换的兼容性都会好很多。另外提醒一下部分推理引擎对 multi-head attention 的算子支持程度不一样。NCNN 在转换时如果报错常常是遇到了一个复杂度较高的矩阵乘法组合。我的处理方案是先看 ONNX 里对应节点的算子类型然后手写一个 custom layer把多层矩阵计算合并成几个基本算子。这个步骤比较繁琐但一次写好后续就能稳定复用。4.2 NCNN 与 TensorRT 实测对比我分别在 NCNN 和 TensorRT 上做了推理测试。TensorRT 对 Transformer 的支持相对好FP16 下 640x640 输入整合模型单帧推理时间大约 7.2msNCNN 在 CPU 端的表现则要看设备我在一颗移动端 SoC 上测出来FP32 模型单帧大概 28ms加上 INT8 量化后能到 18ms。相比之下原版 YOLOv5s 在同设备上 FP32 是 22msINT8 是 15ms。也就是说引入 MobileViT 之后CPU 部署的推理时间增加了 20% 左右但换来的是精度提升这个交换我觉得是划算的。如果你的设备推理能力很有限可以尝试另一种组合保留 P3 阶段的 MobileViTBlockP4 阶段换回 C3 原版这样精度提升会小一点但推理时延能压到基本持平。4.3 量化时保护 LayerNorm 和 SoftmaxMobileViTBlock 里面带 LayerNorm 和 Softmax这两个算子对量化非常不友好尤其是小目标检测场景。直接做全层 INT8 量化mAP50 可能会掉到只剩骨架上的一半。我实测发现把 LayerNorm 留在 FP32只量化其他卷积层精度掉控制在 1.2 到 1.5 个点以内比全量化好得多。所以在量化配置里一定要把包含 LayerNorm 和 Softmax 的节点加入“排除量化层”名单。这个操作在大部分推理框架里都有配置项只是默认不开需要我们手动指定。5. 常见问题与排查技巧实录5.1 训练阶段高频问题速查表问题现象排查思路解决方案训练不收敛loss 在高位震荡学习率偏大、warmup 不够初始学习率降到 0.0025warmup 加到 3-5 epoch显存爆炸前向就 OOM浅层特征图太大序列长度过长不在 P2 阶段引入 MobileViT增大 patch_size减小 batch精度不升反降mobilevit 替换越多越差模型过拟合或序列注意力干扰 anchor去掉 P5 阶段的 MobileViT保留 P3/P4BN 统计异常验证集 mAP 波动大batch size 太小保持 batch 大于 16必要时启用 syncBN权重加载报错state dict 尺寸不匹配新旧结构参数名不一致写一个对齐脚本逐个 key 匹配复制这里面最容易被忽略的是 BN 统计异常。C3MobileViT 里虽然含 LayerNorm但 YOLOv5 里的 Conv 模块仍然带 BN如果 batch size 设成 4 或 8BN 的统计量根本估不准训练过程会剧烈震荡。我建议融合模型至少用 batch16 起步。5.2 部署阶段高频问题速查表问题现象排查思路解决方案ONNX 导出失败遇到不支持的 Reshapeunfold/fold 动态尺寸固定输入尺寸 640x640opset12NCNN 转换失败矩阵乘法组合不被识别multi-head attention 拆解不够手写 custom layer 或降级算力量化后精度崩了小目标全丢LayerNorm/Softmax 被量化把这些节点排除量化名单推理速度反而变慢比原版慢很多推理引擎对 transformer 无优化换 TensorRT 或升级算子库测试不同 patch_size5.3 两个必须提前规划的细节第一MobileViTBlock 在 PyTorch 里和 ONNX 里的行为不完全一致。PyTorch 的 unfold 和 fold 在某些推理引擎中看起来能转过去但推理结果会跟 PyTorch 对不上最常见的原因是 padding 的默认值不同。建议转换完 ONNX 后先用 Python 的 onnxruntime 对同一张输入图做输出对比误差超过千分之一就说明导出环节出了问题提前排查不要等到上板了再查。第二如果想在真实移动端 App 里跑尽量减少 H/W 之外的动态分支。MobileViTBlock 这种 patch 切分操作会天然引入和输入尺寸强相关的逻辑尽量不要在 app 里做任意尺寸的实时推理。比较好的方案是服务端算好 anchor 和预处理参数客户端只跑固定分辨率的模型。6. 可复现准备与资料参考最后把我这次实验用到的主要参考资料整理一下方便想尝试的朋友按图索骥。YOLOv5 官方仓库模型结构定义、训练脚本、anchor 聚类脚本都在里面。MobileViTv1 原始论文重点看通道数配置和 MobileViTBlock 的消融实验。torchvision 中的 MobileViT 实现源码可以对比看 unfold/fold 的标准写法。NCNN 官方算子支持列表转换前先查一下 MultiHeadAttention 相关算子是否被支持。ONNX Simplifier 工具解决冗余节点和动态 Shape 问题的利器。这些资料我在项目里都实际用过边看边改代码效率最高。如果你也想动手复现我建议别急着上来就改一堆结构先原样训练一个 YOLOv5s 的 baseline然后把 MobileViTBlock 加进去跑一版等确认收益后再决定要不要往更深层铺开。我个人在实际操作中最深的体会是这种“混合结构”的坑基本不在训练阶段而在部署阶段。训练阶段大多数问题都可以通过减少替换层、控制序列长度解决但一旦到 ONNX 转换和移动端算子适配你才能感受到轻量 Transformer 模型“看起来轻跑起来不一定轻”的复杂度。先想清楚自己的推理引擎支持什么算子再决定怎么改模型结构远比你先把模型训好了再回来补兼容性要省事得多。希望这篇记录能帮你少走几条弯路。
