简介:基于PyTorch的MobileNet-YOLO检测网络实现,融合mobilenet系列(v1/v2/v3)与yolo系列(yolov3/yolov4)的组合方式,适合希望在嵌入式或轻量级设备上部署目标检测模型的研究者与开发者。项目沿用与原始实现非常接近的损失函数,在VOC2007+2012上完成训练、在VOC2007上测试,预训练权重采用ImageNet而非COCO;输入分辨率覆盖352等档位,便于对照论文实验结果或迁移到自建数据集。压缩包共29个文件,包含16个Python脚本、5个Shell脚本、2个YAML配置、4张示例图像以及说明文档和License文件,整体仅240KB,结构轻量、易于阅读。Python脚本覆盖模型定义、训练、推理、评估、数据读取与图像增强等环节;Shell脚本封装VOC数据下载、LMDB创建、训练和推理流程;YAML配置便于调整数据与模型参数。目前已有1943人学习下载,适合实践轻量化目标检测、熟悉MobileNet+YOLO训练流程或在此基础上二次开发的读者。 最近在做移动端目标检测项目时,我整理了一套基于Pytorch的Mobilenet-YOLO仓库,把MobileNet系列的轻量骨干(v1、v2、v3)和YOLO系列的检测头(yolov3、yolov4)组合到一起,用来快速验证不同backbone在不同业务场景下的表现。这个仓库本身不是什么新网络,而是提供一个“骨架清晰、能随便换零件”的训练与部署框架,正好解决了我这边反复在效率和精度之间来回试错的痛点。如果你也在做边缘端检测、模型轻量化或者想把检测模型落地到嵌入式设备,这篇东西值得看完。
我最早是从YOLOv3开始接触目标检测的,后来为了在树莓派和RK3588这类设备上跑实时检测,又把MobileNet系列翻了个底朝天。慢慢发现,很多网上的教程都只讲YOLO本身,或者只讲MobileNet分类,真正把两块拼在一起、用同一套Pytorch代码去训练和部署的中文资料少之又少。所以我决定把自己踩过的坑、整理过的模块结构、以及实际调参的经验一次性写清楚,希望能帮你少走弯路。
1. 项目缘起:为什么把MobileNet和YOLO放在一个仓库里
1.1 目标检测的两条技术主线
目标检测网络发展到今天,基本可以拆成两个独立的组成块:backbone负责提取特征,head负责在特征图上预测目标的类别和位置。早期像YOLOv3、SSD这些模型,backbone大多是VGG、Darknet这类“重量级选手”,精度不错,但放在嵌入式设备上根本跑不动。后来MobileNet系列出现,用深度可分离卷积大幅减少了参数量和计算量,很快成为轻量化检测模型的首选骨干网络。
另一条主线是YOLO系列的检测头。从YOLOv3的多尺度预测,到YOLOv4引入CSPDarknet、PANet和Mish激活函数,再到后来YOLOv5、v8的anchor-free和解耦头,YOLO一直在解决“怎么把特征图上的信息变成可靠的目标框”这个问题。我在做项目时发现,检测头本质上和后端业务关系不大,真正决定模型能不能跑起来的是backbone。所以把MobileNet和YOLO放在同一个代码仓库里,本质上就是提供一个“自由组合”的试验台:今天用MobileNetV3+SELC检测头,明天换YOLOv4的PANet,只需要改一行配置。
1.2 这个仓库解决了什么
这个仓库解决的最大问题,是避免重复造轮子和降低切换成本。之前我手头有基于Darknet的YOLOv3训练脚本,也有基于MobileNet的分类模型,但两者代码风格完全不同,一个用C语言的darknet,一个用Pytorch,没法直接拼。统一到Pytorch之后,我可以把MobileNetV1/V2/V3的权重直接加载到检测模型里作为初始化,然后只训练检测头部分,训练速度提升明显,而且显存占用也少很多。
另外,仓库里会把训练、验证、导出、推理这几个环节串起来。比如我训练完一个MobileNetV3+YOLOv4模型,直接导出ONNX,再转成TensorRT或者RKNN放到板子上跑,整个过程不需要动网络结构代码。这一点在实际项目中很关键,因为算法团队和部署团队经常要来回对齐,如果训练代码和部署代码是两套,很容易出现精度对不上的问题。统一框架之后,至少能保证训练和导出的计算图是同一个。
2. MobileNet系列与YOLO系列的选型要点
2.1 MobileNet v1/v2/v3的演进逻辑
MobileNetV1的核心是深度可分离卷积,把标准卷积拆成逐通道卷积和1x1逐点卷积。这个操作带来的计算量下降非常明显,但有一个坑:逐通道卷积处理后的特征图直接经过ReLU,容易把低维信息丢掉。MobileNetV2提出的Linear Bottleneck和Inverted Residual正是为了解决这个问题,先用1x1卷积升维,再做深度卷积,最后用1x1卷积降维,并且在降维后不使用ReLU。
MobileNetV3则更进一步,引入了NAS搜索得到的结构、Squeeze-and-Excitation注意力模块和h-swish激活函数。实际用下来,V3在ImageNet分类上比V2精度更高,速度也更快。但在检测任务里,V3的SE模块虽然提升了一定精度,却在某些硬件上(比如树莓派)带来额外的计算开销。我通常会在“追求极致速度”的场景选V1或V2,在“精度和速度平衡”的场景选V3,具体看部署平台的算力。
2.2 YOLOv3到YOLOv4的检测头变化
YOLOv3最大的贡献是引入多尺度预测和FPN(特征金字塔),分别在三个不同尺寸的特征图上做预测,用来解决目标大小不一的问题。YOLOv4在v3的基础上做了三件事:backbone换成CSPDarknet53,neck改成PANet,head保留YOLOv3的anchor-based结构,但训练时加入了Mosaic数据增强、CIoU损失和DIoU NMS。这些改动让YOLOv4的mAP提升了一个档次,但计算量也随之增加。
当我把YOLOv4的neck和head接到MobileNet的骨干上时,有一个很微妙的问题:PANet的上下路径会反复融合特征,MobileNet的特征图本身比较“薄”,信息量不如Darknet丰富,所以检测头不能照搬yolov4默认的channel数量,需要适当减小。我在仓库里加了一个channel_scale参数,默认0.5,能在不损失太多精度的前提下大幅降低计算量。
2.3 组合选型的实际操作
组合选型时,先想清楚三件事:训练样本量、目标大小、部署硬件。如果训练数据只有几千张,我建议用MobileNetV2+YOLOv3这种组合,参数少,不容易过拟合;如果是大规模数据集,且有GPU资源,可以考虑MobileNetV3Large+YOLOv4。如果目标以小物体为主,必须保留多尺度预测,并且把输入分辨率拉到640以上;如果是针对视频监控这种固定场景,可以考虑只保留两个尺度的预测,减掉一个head分支来提速度。
我自己的经验是:先跑通最小可用版本,再逐步加复杂度。比如先用MobileNetV1+YOLOv3的默认配置训练一轮,看loss曲线是否正常,然后再切换成V2/V3,观察提升幅度。如果提升不明显,优先去调数据增强和训练策略,而不是继续换网络。很多刚入门的朋友一上来就追最新模型,结果训练出的结果还不如老老实实调参的YOLOv3,原因就在这儿。
3. 在Pytorch里统一训练与部署的核心实现细节
3.1 网络结构的模块化设计
仓库结构上,我把backbone和head拆成独立模块:MobileNet系列都实现backbone.py里,每个类输出三个尺度的特征图;YOLO系列检测头放在head.py里,接收backbone输出的特征图列表,经过neck和head生成预测。这样当一个backbone改名为yolo backbone时,对检测头完全透明。
具体实现时,backbone的输出通道数要保持和YOLO head的输入通道匹配。MobileNetV1的输出通常取最后几个stride层,比如对输入尺寸416,输出stride分别是8、16、32,对应通道数需要自己规定。YOLOv3的head会对每个尺度的特征图各做一次3x3卷积和1x1卷积,得到预测张量。对于coco的80类,每个尺度输出shape是 (batch, 3*(5+80), H, W),其中3表示每个网格3个anchor。如果用MobileNet的backbone,channel数变小了,检测头的最后一个卷积输出维度不变,所以只要在中间加一个1x1卷积调整通道即可。
这里要注意,不要让“模块化”变成“过度抽象”。我刚开始写的时候,把每个卷积、BN、激活都封装成函数,结果debug时根本看不清楚每一层在干什么。后来改成用一个build_backbone函数,用配置字典初始化网络结构,参数一目了然,也方便出图。
3.2 训练配置与损失函数细节
YOLO系列训练有一个容易忽略的点:除了常规的分类损失和回归损失,还需要对objectness(是否有目标的置信度)做区分。对于正样本,objectness的标签是anchor与gt的IoU;对于负样本,让objectness趋近0。如果直接照搬分类模型的Pytorch训练代码,很可能忘了在损失函数里加上ignore阈值,导致收敛不到最佳点。
仓库里我配置了CIoU loss作为回归损失,因为相比Smooth L1,CIoU能同时优化重叠面积、中心点距离和长宽比,在密集小目标场景下提升比较明显。训练时用SGD优化器,初始学习率0.001,配合余弦退火,batch size设为32。如果显存不够,可以把输入分辨率从416降到320,损失函数不用改,检测精度会略降,但训练速度提升明显。
另外,锚框的选择也很重要。YOLOv3默认的9个anchor是针对COCO数据集聚类的,如果你自己的数据集里目标尺寸差异很大,强烈建议重新用KMeans聚类,不然训练一开始loss就很难降。仓库里有mmdet_utils/kmeans.py,可以直接生成新的锚框配置。
3.3 模型导出与端侧部署
训练完模型之后,导出环节通常是最折腾的。我的推荐流程是:Pytorch权重 -> ONNX -> TensorRT或RKNN,导出时要把一些自定义操作(如NMS、Resize)剥离出去,只保留纯卷积运算部分,这样转换最稳定。
在Pytorch里导出ONNX时,有一个常见问题:如果代码里用了Tensor.size()导致的动态shape,ONNX导出会报错。我一般把输入固定为动态axes,输出只保留三个特征图,后处理NMS放在推理脚本里。转换完成后用onnxruntime跑一遍,对比Pytorch输出的数值差,如果max abs error小于1e-4基本没问题。
对于端侧部署,如果你用的是RK3588,建议直接导出RKNN格式;如果是NVIDIA Jetson,导出TensorRT能获得比较大的加速。我对比过同样是MobileNetV3+YOLOv3,在Jetson Nano上TensorRT FP16相比Pytorch FP32推理速度提升约2.5倍,显存占用只有原来一半,非常划算。
4. 常见问题与工程避坑实录
4.1 环境搭建:Pytorch与CUDA版本
很多人在一开始搭环境就被卡住,尤其是Windows下装GPU版Pytorch。这里我分享一个稳定组合:Pytorch 1.13 + CUDA 11.7,如果不想折腾,直接用Pytorch官方命令安装。注意先查自己显卡驱动支持的CUDA版本,用nvidia-smi看右上角Driver版本对应的CUDA Version,只要驱动版本 >= 你安装的CUDA runtime版本即可。
另外,2024年之后Pytorch更新很频繁,新版本默认不再支持老显卡。如果你显卡是GTX 10系,建议装Pytorch 1.x版本;如果是RTX 30/40系,用Pytorch 2.x配合CUDA 11.8。我在一台只支持CUDA 11.8的电脑上安装Pytorch时,最开始用默认源下载很慢,改用手机热点结果更慢。解决办法是换阿里云镜像源,或者下载本地whl文件安装。总结一句:先确定CUDA版本,再确定Pytorch版本,最后再想模型怎么选。
4.2 参数量前后不一致的现象
有朋友遇到过一个问题:YOLO训练一开始统计的参数量和最后训练完得到的参数量不一致。这通常不是bug,而是因为训练过程中使用的EMA(指数移动平均)权重没有计入统计,或者类别数配置在训练中途被修改。如果你是先加载预训练模型再继续训练,加载的权重结构可能和当前配置不完全一致。
另外,如果你在训练时开启了DDP(分布式数据并行),模型参数量在初始化阶段只统计了一部分,会有一个all_reduce过程让每个进程看到完整参数。如果统计代码写在模型构造之后、DDP包装之前,就会漏掉一部分参数。我建议统一在model_ema对象构建完成后统计参数量,这样前后就一致了。
4.3 重叠框与NMS处理
目标检测里“模型输出一堆重叠框”是再正常不过的事,因为你会在每个网格上预测多个anchor。后处理时用NMS去掉重叠度高的框。YOLOv4里用的DIoU NMS比普通NMS更适合密集场景,因为它在IoU基础上加入了中心点距离惩罚,两个形状完全一样但中心点接近的框会更不容易被误删。
如果你发现NMS之后还是有很多重复框,通常是置信度阈值设太低(比如0.1),或者类别数设置错误,导致所有预测都归为同一个类别。另一个常见的坑:CPU上跑多进程推理时,NMS处理速度反而很慢,因为每个进程都要独自load模型和预处理。我之前测试,在CPU上开4个进程处理一段视频,总耗时比单进程还慢1.4秒,原因就是CPU多进程通信开销远大于计算时间。解决办法是用OpenCV的并行方案或者TBB,而不是简单的multiprocessing。
4.4 导出模型给QT调用
很多桌面应用想接YOLO检测,但训练完模型不知道怎么转成能被QT调用的格式。我推荐两种方式:一是把模型导出为ONNX,再用OpenCV的DNN模块读入。QT里调用OpenCV DNN比调用Pytorch轻量得多,而且不需要额外安装Python环境;二是直接把Pytorch模型打包成libtorch库,通过C++ API调用,但要处理Tensor和OpenCV Mat的转换,稍烦琐。
如果模型里用了自定义激活函数(比如h-swish),导出到ONNX时ONNX Runtime不一定支持,必须在导出前把h-swish替换成标准的ReLU6或者Relu。MobileNetV3的h-swish在ONNX里通常可以表示为hard_sigmoid乘x,但某些版本会报错。我踩过这个坑,后来在导出函数里加了inplace替换,把h-swish换成普通swish,部署时精度下降可忽略,但兼容性好很多。QT调用时,把后处理NMS用C++重写,加上一个简单的线程池,就能实现60FPS的实时检测窗口。
说到最后,真正把MobileNet和YOLO拼在一起跑通,其实并不复杂,难的是每一步都有人告诉你“这里会踩坑”。我在这套仓库里最满意的不是某个网络结构,而是把所有项目里遇到过的环境、训练、导出问题都沉淀成了排查清单。下次你换了新板子或者新数据集,照着清单走一遍,基本都能顺下来。如果你也正在折腾轻量化检测,建议先从小网络开始跑通整个流程,再逐步升级,别一上来就追最新模型——工程上的稳定,永远比模型上的“新”更值钱。
本文还有配套的精品资源,点击获取