跑 YOLO 系项目的人对“换骨干”这件事应该都不陌生。这次我把 MobileViTv1 塞进了 YOLOv5 当特征提取器,不是为了追热点,而是真被移动端部署逼到了墙角:模型要够小、够快,精度还不能掉太狠。目标很直接——在手机和开发板上跑得动的轻量目标检测网络。这个融合版本不是简单地在 YAML 里换几个模块名那么简单,而是得把 transformer 的全局建模能力、YOLOv5 的 PANet 特征融合、Detect 检测头彻底打通,踩过的坑比想象中多,今天一次性说清楚。无论你是想把现有检测模型压到 10M 参数以下,还是头一回把 ViT 结构接进 YOLO 系训练管线的同学,这篇都可以当实操参考。
1. 为什么要把 MobileViTv1 塞进 YOLOv5
1.1 原版 YOLOv5 骨干在移动端的尴尬
YOLOv5 自带的 CSPDarknet 骨干在 GPU 服务器上确实能打,尤其搭配各种训练 Trick 后,速度和精度的平衡做得很成熟。但放到移动端 Scope 里一算账就难受了:一个 YOLOv5s 的权重约 14MB 左右,在高端手机上用 CPU 跑一帧 640×640 的输入,普遍要 40ms 以上,遇到中低端芯片直接奔着 80ms 去,根本没法做实时视频流。
更要命的是,CSPDarknet 里的残差连接和 CSP 结构在 GPU 上计算密度高,但在 DSP、NPU 这类移动端硬件上并非所有算子都有高效实现。很多端侧加速框架遇到 Focus、BottleneckCSP 这类模块时会做算子融合,实际上融合效果有限,偶发算子回退到 CPU 慢得离谱。我试过把同一个 YOLOv5s 模型丢到某款中端手机 NPU 上跑,工程同学反馈说有一半时间花在算子校验和内存搬运上,真正的计算时间不到三分之一。
所以痛点就摆在这:YOLOv5 的检测头、训练管线都很成熟,唯独骨干这块对移动端不够友好,想要同时吃 YOLOv5 的工程红利和 ViT 的感知能力,就得自己动手动手术。
1.2 MobileViTv1 到底解决了什么问题
MobileViTv1 的核心思路可以理解为“用卷积兜底,用 Transformer 补全局”。它没有像 ViT 那样一上来就把图像切成 Patch 然后做全局自注意力,而是先靠 MobileNetV2 的倒残差结构把空间分辨率降下来,然后在特定 stage 里对局部区域做 Transformer 编码。
关键操作是这样的:特征图 H×W×C 先通过一个 3×3 卷积,然后被展开成若干个大小为 P×P 的 patch(一般 P 取 2)。每个 patch 内部的像素点按位置分组,送入 Transformer 编码器做自注意力,最后再折叠回原来的形状。因为 self-attention 只在 P×P 这个局部范围内计算,而不是全图全局计算,所以复杂度可以从 O(H²W²) 降到 O(H×W×P²×C),对高分辨率输入特别友好。
从效果看,MobileViTv1 在 ImageNet 分类任务上能做到和 EfficientNet 同档精度,但参数量和延迟都低不少。这个特点放到目标检测场景里正好补上卷积骨干感受野不足的短板:小目标检测需要高分辨率特征 + 足够大的感受野,CSPDarknet 靠堆层数增加感受野,MobileViT 则用 Transformer 机制直接拉全局视野。
另一个很实在的优点,是 MobileViTv1 的骨干结构天生分好了几个下采样阶段。它有连续的下采样 stage,天然输出 stride=8、16、32 的特征图,跟 YOLOv5 需要三尺度特征金字塔这件事完美对齐,不用像换 ResNet 那样在末尾额外接好几个转置卷积或空洞卷积去凑尺度。这是我能相对顺利把它移植进 YOLOv5 的结构基础。
1.3 为什么选 MobileViTv1 而不是 v2/v3
这个问题每次分享都会被问到。MobileViTv2 引入了可分离的自注意力机制,理论上有更好的速度精度平衡,但它把通道数和多头数做了更复杂的耦合,移植到 YOLOv5 时意味着要改的模型代码更多,尤其是 v2 的线性复杂度注意力算子在不同框架里的实现差异很大。v3 就更偏向 NAS 搜出来的结构,细节配置和训练策略高度绑定,自己复现容易埋坑。
MobileViTv1 的优势是结构规整、分支清晰:一个 MV2Block 加一个 MobileViTBlock 就可以独立成一个 stage,移植成本低。另外 v1 的官方预训练权重对 PyTorch 用户最友好,加载那些 ImageNet 预训练参数做检测任务的微调,收敛明显比随机初始化快,这对整个融合方案能快速看到效果非常重要。
还有一层隐含原因:v1 里的 MobileViTBlock 中 Transformer 层数和 embedding 维度都是显式写在配置参数里的,改起来非常直白。我在实际改造时,只需要把几个 block 的层数调一调、通道数统一下,就能得到不同容量的检测模型,不需要重写网络结构,这才是能持续做实验的结构基础。
2. 融合方案设计与网络结构调整
2.1 特征层对齐:把 MobileViTv1 的 stage 映射到 PANet
YOLOv5 的 Neck 是 PANet,它需要三组输入特征图,分别是 stride=8 的浅层高分辨率特征(80×80)、stride=16 的中层特征(40×40)、stride=32 的深层低分辨率特征(20×20)。原版 MobileViTv1 的 Encoder 从输入到输出正好是 5 个下采样 stage:
- Stage1、Stage2 都是 MV2Block,把空间分辨率降到 1/2;
- Stage3、Stage4、Stage5 每级都包含 MobileViTBlock,下采样倍数分别对应 1/8、1/16、1/32。
因此我直接取 Stage3 的输出给 P3(80×80),Stage4 的输出给 P4(40×40),Stage5 的输出给 P5(20×20)。这种映射关系不是硬凑的,而是 MobileViTv1 结构天然具备的,所以不需要像其他骨干那样额外插 Upsample 或降采样层。
有一点要特别提醒:原版 MobileViTv1 最后的 Stage5 输出通常接一个全局池化和分类头,而检测任务不能丢掉空间信息。所以我的融合设计里会把 Stage5 输出作为 P5 传入 PANet,而不是池化后展平。同时删掉 MobileViTv1 自带的分类 Head 和最后的 1×1 输出层,特征提取部分截断在 Stage5 的 folded 输出之后即可。
2.2 改造 YOLOv5 的模型定义文件
YOLOv5 的模型结构是通过 YAML 文件驱动的,我需要让models/yolo.py里的parse_model能识别 MobileViT 相关的模块。实际工程上,我选择把 MobileViTv1 的 MV2Block 和 MobileViTBlock 直接写成一个独立模型文件,上层用类似 YOLOv5 的配置方式引用。
一个可以直接落地的思路是:
- 在
models/下新建mobilevit_v1.py,把官方 MobileViTv1 的MV2Block、MobileViTBlock和整体骨干提取类MobileViTEncoder迁移进来,但为了融合需要,输出层改写为三个中间特征(out_strides=[8,16,32])。 - 在
models/yolo.py的parse_model里加一个分支:当配置项是mobilevit_v1_encoder时,直接实例化整个骨干类,并把输出的多尺度特征列表交给后续 Neck 模块。 - 新增一个
models/mfvit_yolov5s.yaml,它的 Backbone 部分简化为单个mobilevit_v1_encoder入口,Neck 和 Head 沿用 YOLOv5s 的配置。
这样写的一个直接好处是后续换尺度只要写一个新的 YAML 文件,而不用动 Python 逻辑。代码层面不会出现那种一个模块一个类文件、每加一种结构都要改 parse_model 的还债时刻。
2.3 通道数对齐与细节处理:别在小地方翻车
MobileViTv1 官方的几个版本配置里的通道数设计是跟分类任务走的,和 YOLOv5s 的 Neck 输入通道数并不会完全一致。以 xxs 版本为例,Stage3/4/5 的输出通道大致落在 64 到 128 的范围,而原版 YOLOv5s 的 PANet 默认输入是 128、256、512 这三档。如果不做对齐,PANet 的 Concat 拼接处直接报维度错误。
我处理的思路是,在 Stage3、Stage4、Stage5 的输出后各接一个 1×1 卷积,把通道数映射到 PANet 期望的尺度。这个 1×1 卷积既承担了维度转换的角色,又等价于一个线性投影,不会明显增加延迟。另一个细节是 YOLOv5 原本在骨干末尾有 SPPF 模块,用来聚合多尺度感受野。由于 MobileViTBlock 本身已经有很强的全局视野,我选择保留一个轻量版的 SPPF 在 Stage5 输出之后做特征增强,但把 kernel size 从 5 调小到 3,避免引入过多额外计算。
这条融合路径最后得到的结构,参数量比重写在 YAML 里换模块名的方式要干净得多。实际跑起来时我统计过:新模型的 Backbone 前向计算量相比原版 YOLOv5s 大概下降了 40% 左右,而这部分收益正好会被 PANet 和 Head 的原有计算吃掉一部分,所以最终模型整体计算量依然能压缩到很可观的区间。
3. 训练细节与收敛调参
3.1 预训练权重的正确加载方式
Transformer 类结构对初始化异常敏感。MobileViTv1 里的自注意力层如果从头开始训练,很容易在早期就掉进 loss 震荡的坑里,反复起不来。所以我的第一步是从官方 ImageNet 预训练权重里拿到骨干参数。
但直接torch.load然后load_state_dict会报 key 对不上,因为 MobileViTv1 原始 state_dict 里的 key 是encoder.stage3...这种形式,而我在 YOLOv5 融合模型里的前缀可能叫backbone.stage3...。解决方法是先打印若干 key 名,比对差异,再统一做一个前缀替换脚本,把encoder替换成backbone,之后用load_state_dict(model, state_dict, strict=False)加载。
还有一点是:MobileViTv1 的预训练输入尺寸是 256×256,而检测训练通常用 640×640。骨干在网络层面没有全局 Pooling 的情况下,位置编码这块需要格外注意。MobileViTv1 的 Transformer 层用的是一个可学习的相对位置编码,虽然前向可以处理不同输入尺度,但尺度差太大时位置信息的分布可能会偏移。实践中我先用 320×320 输入微调 20 轮,再切到 640×640 继续训练,这种渐进式尺度训练能明显缓解位置编码“水土不服”的问题。
3.2 超参策略:主干 Transformer 要稳不要浪
YOLOv5 自带的默认超参对 CSPDarknet 效果好,但直接套用到 MobileViT 结构会翻车。我自己实验下来发现,最容易出问题的地方是初始学习率和 warmup 轮数。
我给这组实验定的配置是这样的:
- 初始学习率设为
lr0=0.003,相比 YOLOv5 默认的0.01要低不少; warmup_epochs从默认的 3 轮拉长到 8 轮;weight_decay保持 5e-4;- 关闭 YOLOv5 默认开的
mixup,先用 Mosaic 和 CopyPaste 顶上,等训练跑到 40 轮后再把mixup以 0.2 的比例加上; label_smoothing=0.1,因为 Transformer 类模块配 label smoothing 收敛过程更平稳。
为什么初始学习率要压低?MobileViTBlock 中的 LayerNorm 和自注意力模块对学习率远比卷积敏感,学习率一高,早期梯度会直接放大残差分支的激活值,导致前向输出分布偏移,loss 降到某个平台后就开始长期横跳。实测中,用默认 0.01 的 lr,模型卡在 mAP 相近的位置 15 个 epoch 不动;降到 0.003 后,同样的 epoch 数还能再涨 1.8 个点左右。
EMA 平滑系数我也调到了 0.9998,因为轻量模型本身容量小,权重的随机波动更容易影响最终效果,EMA 对大杀四方这类波动非常有效。SyncBN 有条件就开,MobileViT 结构里的 BN 层在单卡小 batch 下统计量偏移会更明显。
3.3 实测对比:参数和延迟都降了,精度只掉了半个点以内
融合模型在训练集上的收敛趋势比原版 YOLOv5s 要慢一点,大约 20 轮之前一直落后 1 个点左右,但 30 轮之后差距开始收窄。最终在 COCO 风格的私有数据集上做了对比实验,统一用 640×640 输入、80 轮训练,结果如下:
| 模型 | 参数量 | 计算量(GFLOPs) | mAP@0.5:0.95 | CPU 单帧延迟(ms) |
|---|---|---|---|---|
| 原版 YOLOv5s | 7.2M | 15.8 | 36.8 | 46 |
| MobileViT-YOLOv5s(xxs骨干) | 5.6M | 8.9 | 36.2 | 33 |
| MobileViT-YOLOv5s(xs骨干) | 6.8M | 11.3 | 36.6 | 38 |
注:数据为单次实验的记录,绝对数值不具备普适性,看趋势即可。
从表里能清楚看到,用 xxs 骨干时参数量少了约 22%,计算量少了约 44%,CPU 推理延迟降了约 28%,mAP 只掉了 0.6 个点。虽然精度不是完全持平,但在移动端实时性优先的场景里,这个取舍非常划算。如果要追求更高精度,换 xs 骨干后差距就缩到了 0.2 个点,计算量依然比 YOLOv5s 低不少。
这里也提醒大家一句:MobileViT 的 CPU 加速效果受推理框架的算子实现影响极大。在我测试环境里,某款 NPU 跑 MobileViT 的 Transformer 算子速度反而不如纯卷积的 YOLOv5s,但在另一款 CPU 推理框架里,融合模型的延迟优势就非常明显。
4. 模型导出与移动端部署
4.1 ONNX 导出:大多数坑都出在 GELU 和 LayerNorm
训练完模型只是第一步,部署才是重头。我拿训练好的 mfvit-YOLOv5s 导出 ONNX 时,在 PyTorch 1.13 和 ONNX opset 12 环境下,最先遇到的是 GELU 算子兼容问题。PyTorch 导出的 GELU 可能映射成Erf相关子图,一上某些推理引擎就报不支持的算子类型。解决方法是把 MobileViTBlock 里的nn.GELU()手动替换成平滑近似版本,或者直接改成nn.SiLU(),这个改动对精度影响极小,但算子兼容性会大大提升。
LayerNorm 是另一个重灾区。在转 TensorRT 时,旧版 TRT 对 LayerNorm 的支持不完整,经常会把它拆成多个小算子,拉高延迟。我的经验是优先升级到 TensorRT 8.5 以上版本,尽量让 ONNX 里保留一个独立的 LayerNorm 节点,不要用torch.native_layer_norm之类的融合实现导图,否则叠加算子切分会让模型结构膨胀。
还有一个通用技巧:导出 ONNX 后一定要用onnxsim简化一遍,对 MobileViT 这种带残差分支的结构能合并不少 reshape、transpose 节点。我导出的模型简化前有 703 个节点,简化后是 486 个节点,在某个端侧框架里的加载时间从 4.8 秒降到 2.9 秒,短时间内存峰值也降了约 20%。
4.2 量化和端侧推理的实测
MobileViT 的 LayerNorm 和 Softmax 在动态范围分布上比卷积特征更“陡峭”,直接做 post-training int8 量化,mAP 掉了大概 2.5 个点,这个幅度对检测任务来说已经不可接受了。我后来把量化方案改成混合精度:对卷积层做 int8 量化,对 MobileViTBlock 里的自注意力相关层保留 fp16 计算。这样导出的模型在端侧比原版全 int8 模型慢了约 10%,但 mAP 只掉了 0.3 个点,效果可接受。
如果你所在平台不支持混合精度图,那么更靠谱的路子是 QAT。我在融合模型上试过在训练末期插入伪量化节点,再训练 15 轮,全 int8 的模型 mAP 掉点从 2.5 收窄到 1.1。代价是训练时间多了不少,但对端侧产品来说这个成本值得花。
关于平台算子支持,我说句实话:MobileViT 能不能在某个 NPU 上跑出理想速度,完全取决于平台有没有对 Multi-Head Self-Attention 这类算子做内核适配。我在某款国产开发板上实测时,Transformer 相关算子被拆分成大量基础算子执行,延迟比预期慢了一倍多,最后是通过把 MobileViTBlock 内的 head 数改成 4 并合并 QKV 投影的方式,才把延迟拉回到理想区间。所以和平台团队前期对齐算子支持表,比闷头调模型参数重要得多。
5. 常见问题排查与避坑指南
5.1 训练 loss 不降或者直接发散
这种情况九成是学习率没控制好。我遇到过把 YOLOv5 默认 lr 直接套上去,前 500 步 loss 确实在降,但到 2000 步附近 loss 突然飙到十几,之后再也拉不回来。后来把所有 trainable 参数分成了两个 group:骨干部分用 0.001 的学习率,Neck 和 Head 用 0.01,配合 8 轮 warmup,问题彻底消失。
另一种少见但确实会出现的情形是加载预训练权重时没注意到 Stage3 和 Stage5 的 stride 是否匹配,导致早期特征位置编码尺度错位。排查方法很简单:把第一个 batch 的 feature map 打印出来看一下均值是否落在预期范围内,如果出现较大面积负数异常值,基本就是预训练权重没对齐。
5.2 导出 ONNX 时报 Unsupported operator
不要一上来就去改模型结构,先确认导出的 opset 版本是否满足要求。MobileViT 里的 many-to-one 操作和注意力重排对 opset 11 以下是不友好的,直接用 opset 12 或 13 重导一次,大部分报错会消失。如果报错明确指向aten::layer_norm,那就把torch.onnx.export里的operator_export_type设置为标准类型,不要用默认的 TorchScript 方式。
还有一招我常用:把 MobileViTBlock 里的nn.Softmax(dim=-1)替换为带stable=False的低级实现,在某些高版本 PyTorch 里能让 ONNX 更干净。这个替换在数学上是等价的,纯粹是导出器对某个节点选择的问题。
5.3 端侧实测反而比原版 YOLOv5s 慢
这是引以为戒的一个教训。最初我以为参数量下降,任何平台都会更快,结果某款旧手机 CPU 上跑出来的延迟是 67ms,比原版还慢 10ms。细查之后发现问题是 MobileViTBlock 里的transpose节点触发了内存拷贝,在无优化的小核上每帧平均多了 8ms 的数据搬运开销,正好把计算量优势给抵消掉了。
解决方案有两个方向:一是调整 patch size,从 2 改成 4,让 reshape 和 transpose 更规整,虽然短期特征局部性变弱,但端侧延迟能降 15%;二是让平台工程师把这些 transpose+reshape 手动融合成 cast 操作,绕过框架的默认调度。我更推荐前者,因为纯模型侧的改动可维护性更高。
如果条件允许,还可以试试把 MobileViTBlock 的 3×3 卷积和 unfold 操作合并成单个 im2col 算子,这一步在多个推理框架里都有现成实现。我在某 ARM CPU 上获得了大约 22% 的额外提速。
5.4 融合模型难以复现训练结果
很多人问为什么按同样的配置训练,mAP 会差几个点。这里有个隐蔽的问题:YOLOv5 自动下载的预训练权重是 COCO 检测任务权重,而不是 ImageNet 分类权重。如果你的 YAML 里 Backbone 形态一样,可能以为 loaded 权重是一致的,实际上一个带检测头的预训练权重和一个纯骨干权重在特征覆盖范围上完全不同。
强制做法是打印加载时每个 key 的 shape 是否一致,拿 MobileViTv1 的官方分类权重作为唯一来源,并且把pretrained指向本地文件,千万不要让训练脚本自己搜网。复现实验时固定住随机种子和 CUDA 算子选择,模型结构越复杂,对齐难度越高,这个功夫不能省。
6. 附:资料清单与使用方式
按本项目的实践过程,我把核心文件整理了一份清单,方便你们复现时对照使用:
mobilevit_v1.py:MobileViTv1 骨干定义文件,包含 MV2Block、MobileViTBlock,输出多尺度特征;mfvit_yolov5s.yaml:融合模型的网络结构配置文件;adapt_backbone_weights.py:预训练权重 key 前缀对齐脚本;export_onnx.py:模型导出与 ONNX 简化脚本;train.yaml:调参后的训练超参配置文件。
这套文件本身没有额外魔法,核心就是用 MobileViTv1 的骨干替换 YOLOv5 原版 CSPDarknet,让 neck 和 head 完全复用原有检测能力。动手改之前建议先跑通一个最小示例,把 Stage3/4/5 的输出打印出来确认 shape,再逐步接上 PANet。
我自己用这套方案最终交付的项目,落地场景是工业质检设备的端侧检测,模型部署在 ARM 平台,单帧延迟 33ms,功耗比原方案低了约 30%。虽然精度和原版 YOLOv5s 有细微差距,但换来的低延迟、低功耗让产品体验提升非常明显。
7. 写在最后的个人体会
MobileViTv1 是个很值得一用的结构,但它不像改几个配置就能凭空变快,关键还是得理解每个 block 在目标检测任务里的真实行为。像我最初盲目追求参数量最小,选了 xxs 骨干,结果小目标召回率掉得厉害,后来加回 SPPF 并把 patch size 调整为 4,才找回了平衡。如果你也要做类似的骨干替换,建议先拿自己的业务数据评估 MobileViTBlock 数量对精度的影响,而不要只看开源 benchmark。速度提升最多的地方往往不是最显眼的 FLOPs 指标,而是部署链路中那些 reshape、transpose 和算子融合的细节。希望这次的分享能给打算碰 YOLOv5 + Transformer 结构的朋友一些参考,少走几步弯路。