十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Pytorch实战:MobileNet与YOLO组合的轻量目标检测框架

Pytorch实战:MobileNet与YOLO组合的轻量目标检测框架 简介基于PyTorch的MobileNet-YOLO检测网络实现融合mobilenet系列v1/v2/v3与yolo系列yolov3/yolov4的组合方式适合希望在嵌入式或轻量级设备上部署目标检测模型的研究者与开发者。项目沿用与原始实现非常接近的损失函数在VOC20072012上完成训练、在VOC2007上测试预训练权重采用ImageNet而非COCO输入分辨率覆盖352等档位便于对照论文实验结果或迁移到自建数据集。压缩包共29个文件包含16个Python脚本、5个Shell脚本、2个YAML配置、4张示例图像以及说明文档和License文件整体仅240KB结构轻量、易于阅读。Python脚本覆盖模型定义、训练、推理、评估、数据读取与图像增强等环节Shell脚本封装VOC数据下载、LMDB创建、训练和推理流程YAML配置便于调整数据与模型参数。目前已有1943人学习下载适合实践轻量化目标检测、熟悉MobileNetYOLO训练流程或在此基础上二次开发的读者。 最近在做移动端目标检测项目时我整理了一套基于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放在同一个代码仓库里本质上就是提供一个“自由组合”的试验台今天用MobileNetV3SELC检测头明天换YOLOv4的PANet只需要改一行配置。1.2 这个仓库解决了什么这个仓库解决的最大问题是避免重复造轮子和降低切换成本。之前我手头有基于Darknet的YOLOv3训练脚本也有基于MobileNet的分类模型但两者代码风格完全不同一个用C语言的darknet一个用Pytorch没法直接拼。统一到Pytorch之后我可以把MobileNetV1/V2/V3的权重直接加载到检测模型里作为初始化然后只训练检测头部分训练速度提升明显而且显存占用也少很多。另外仓库里会把训练、验证、导出、推理这几个环节串起来。比如我训练完一个MobileNetV3YOLOv4模型直接导出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换成CSPDarknet53neck改成PANethead保留YOLOv3的anchor-based结构但训练时加入了Mosaic数据增强、CIoU损失和DIoU NMS。这些改动让YOLOv4的mAP提升了一个档次但计算量也随之增加。当我把YOLOv4的neck和head接到MobileNet的骨干上时有一个很微妙的问题PANet的上下路径会反复融合特征MobileNet的特征图本身比较“薄”信息量不如Darknet丰富所以检测头不能照搬yolov4默认的channel数量需要适当减小。我在仓库里加了一个channel_scale参数默认0.5能在不损失太多精度的前提下大幅降低计算量。2.3 组合选型的实际操作组合选型时先想清楚三件事训练样本量、目标大小、部署硬件。如果训练数据只有几千张我建议用MobileNetV2YOLOv3这种组合参数少不容易过拟合如果是大规模数据集且有GPU资源可以考虑MobileNetV3LargeYOLOv4。如果目标以小物体为主必须保留多尺度预测并且把输入分辨率拉到640以上如果是针对视频监控这种固定场景可以考虑只保留两个尺度的预测减掉一个head分支来提速度。我自己的经验是先跑通最小可用版本再逐步加复杂度。比如先用MobileNetV1YOLOv3的默认配置训练一轮看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*(580), H, W)其中3表示每个网格3个anchor。如果用MobileNet的backbonechannel数变小了检测头的最后一个卷积输出维度不变所以只要在中间加一个1x1卷积调整通道即可。这里要注意不要让“模块化”变成“过度抽象”。我刚开始写的时候把每个卷积、BN、激活都封装成函数结果debug时根本看不清楚每一层在干什么。后来改成用一个build_backbone函数用配置字典初始化网络结构参数一目了然也方便出图。3.2 训练配置与损失函数细节YOLO系列训练有一个容易忽略的点除了常规的分类损失和回归损失还需要对objectness是否有目标的置信度做区分。对于正样本objectness的标签是anchor与gt的IoU对于负样本让objectness趋近0。如果直接照搬分类模型的Pytorch训练代码很可能忘了在损失函数里加上ignore阈值导致收敛不到最佳点。仓库里我配置了CIoU loss作为回归损失因为相比Smooth L1CIoU能同时优化重叠面积、中心点距离和长宽比在密集小目标场景下提升比较明显。训练时用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()导致的动态shapeONNX导出会报错。我一般把输入固定为动态axes输出只保留三个特征图后处理NMS放在推理脚本里。转换完成后用onnxruntime跑一遍对比Pytorch输出的数值差如果max abs error小于1e-4基本没问题。对于端侧部署如果你用的是RK3588建议直接导出RKNN格式如果是NVIDIA Jetson导出TensorRT能获得比较大的加速。我对比过同样是MobileNetV3YOLOv3在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拼在一起跑通其实并不复杂难的是每一步都有人告诉你“这里会踩坑”。我在这套仓库里最满意的不是某个网络结构而是把所有项目里遇到过的环境、训练、导出问题都沉淀成了排查清单。下次你换了新板子或者新数据集照着清单走一遍基本都能顺下来。如果你也正在折腾轻量化检测建议先从小网络开始跑通整个流程再逐步升级别一上来就追最新模型——工程上的稳定永远比模型上的“新”更值钱。本文还有配套的精品资源点击获取
返回列表