基于Ultralytics/YOLOv5/v8的多任务模型概述本项目基于Ultralytics/YOLOv5通过增加少量计算和显存开销实现了同时进行目标检测和语义分割的多任务模型。在1024×512的输入分辨率下显存增加约350MB而同分辨率下增加一个BiSeNet则需要约1.3GB。此外两个单任务模型独立输入还会带来额外的延迟。模型在Cityscapes语义分割数据集和由Cityscapes实例分割标签转换而来的目标检测数据集上同时训练。实验结果表明检测结果略优于原版单任务的YOLOv5仅限于此实验数据集。分割指标如下s模型验证集mIoU 0.73测试集mIoU 0.715m模型验证集mIoU 0.75测试集mIoU 0.735由于个人将继续考研仓库在tag 2.0发布后近期不会频繁更新issue大概率不会回复问题请参考以下文档震荡爆炸请尝试降低学习率。若后续考上研可能会在本仓库跟进YOLOv5的更新并尝试新的分割结构、增加更多任务、多卡训练和部署的支持。未来版本可能由其他人整理/重构并发布在MANA AI。环境配置和数据集准备环境配置建议直接使用主分支BS2021近期不会再更新功能和结构但如果issue提出BUG空闲时会尽力在主分支修复。$ python-mpipinstall-rrequirements.txt $ python-mpip uninstall wandb注意当前代码未做多卡训练和wandb支持不卸载wandb训练时大概率会有bug。数据集准备本模型的检测和分割数据集独立加载不要求类别相同。当前支持Cityscapes语义分割数据集和实例分割标签生成的目标检测数据集新增BDD100k混合Cityscapes训练支持将BDD100k当作Cityscapes的一个城市组织数据集格式。扩展语义分割数据集需要增改代码继承BaseDataset类。目标检测数据集可以同原版YOLOv5自行替换参考原版YOLOv5文档和./data/cityscapes_det.yaml文件。下载数据自行到官网下载Cityscapes数据集将leftImg8bit和gtFine放入./data/citys文件夹中也可使用citys中的bash脚本下载需要先将脚本中的帐号密码改成你自己的。数据预处理进入citys文件夹下运行citys文件夹的2yolo_filter.py生成目标检测标签。在citys文件夹中新建一个detdata文件夹将生成的images和labels文件夹剪切到detdata中。$cd./data/citys $ python 2yolo_filter.py $mkdirdetdata $mv./images ./detdata $mv./labels ./detdata注比较费空间只做Cityscapes实验的话可以考虑删除复制的图像将leftImg8bit的图软链接到detdata中但不要移动leftImg8bit和gtFine文件夹分割要用。预训练模型2.0版本准备了4种分割Head的预训练模型。从可视化上更推荐PSP和Lab感受野更大速度上没有太大差别Base最快PSP第二Lab和BiSe接近。Base.pt基础版本的分割Head。16层PAN1/8输入配置文件通道512。C3通道略拓宽版本的C3SPPDropout(0.1)1×1卷积到类别。速度精度综合效果不错但SPP配1/8图感受野不够大s模型效果较好但m模型加深加宽后提升不明显。BiSe.pt模仿BiSeNetV1的分割头精度略高速度与Base相似。16,19,22PAN的1/8,1/16,1/32输入配置文件通道无效。ARM改成RFB2增强非线性BiSeNet每个Upsample后有一个3×3卷积refine这里省计算放在Upsample前。Lab.pt模仿DeepLabV3的分割头验证集精度与PSP和BiSe接近速度略慢于PSP和Base与BiSe相仿。4或3,19浅层1/8,PAN的1/16输入配置文件通道256。1/8图1×1卷积到48通道1/16图过RFB1ASPP类似的替代。PSP.pt模仿PSPNet的分割头目前精度最高速度仅次于Base。16,19,22三层融合输入未找到合适的地方放辅助损失放弃辅助损失。预训练模型百度网盘提取码cjxg推理普通图片推理$ python detect.py--weights./pspv5s.pt或其他模型--sourcedata/images--conf0.25--img-size1024data/images里的几张图片来自Cityscapes测试集、BDD100k、Apollo Scape和YOLO。此外data/test_imgs放了一些Apollo的图可以看看Cityscapes训练后在Apollo上的效果使用了BDD100k训练的效果会再好一点。结果图片在runs/detect文件夹中也可以推理同时显示。$ python detect.py--weights./pspv5s.pt或其他模型--sourcedata/images--conf0.25--img-size1024--view-img同尺寸连续帧图片制作视频$ python detect.py--weights./pspv5s.pt或其他模型--source图片文件夹--conf0.25--img-size1024--save-as-video只支持同尺寸图片制作视频例如Cityscapes提供的三段连续帧测试图片Bilibili的demo视频就是这些图。将图片放入同一文件夹内注意若存在不同尺寸图片则结果视频会保存失败。向Cityscapes提交测试集结果$ python detect.py--weights./pspv5s.pt或其他模型--source图片文件夹--conf0.25--img-size1024--submit--no-save开启--no-save不保存结果会快很多并省空间。将测试集6个文件夹的图片合并在一个文件夹进行推理结束后会在此次的runs/detect/此次exp里找到一个results文件夹里面是将trainid转换为id的结果压缩上传官网即可。测速测速使用提交的参数在同尺寸图片上测或推理视频测。图片推理没有开启cudnn.benchmark推理视频文件时或开启--submit或--save-as-video时会开启cudnn.benchmark此时是真正的运行速度。注cudnn.benchmark开启后第一帧图片会测试各种cudnn算子并记录之后同尺寸的每帧图像都会使用最快的算子。cudnn.benchmark仅在同尺寸图像推理时开启否则每输入新尺寸图就会测一遍。训练后测试模型$ python test.py--datacityscapes_det.yaml--segdata./data/citys--weights./pspv5s.pt --img-size1024--base-size1024对比原版多两个参数--segdata后写Cityscapes数据集的文件夹地址现在只支持这个可以参考SegmentationDataset.py自行扩展。检测长边和分割长边参数分离--img-size是检测长边--base-size是分割长边。如果训练后测试你自己的数据集使用test_custom.py训练中train_custom.py会测。$ python test_custom.py--data你的.yaml--segdata你的分割数据路径--weights./pspv5s.pt --img-size1024--base-size1024训练复现训练前先下载对应的原版注意我是在tag V5.0代码上改的COCO预训练模型做初始化见原版README和weights/download_weights.sh脚本。$ python train.py--datacityscapes_det.yaml--cfgyolov5s_city_seg.yaml --batch-size18--epochs200--weights./yolov5s.pt--workers8--label-smoothing0.1--img-size832--noautoanchor不一定如示例训200轮这是我训上述预训练模型为了让其尽量收敛的参数建议最少训80轮我一般训150到180轮。注意为了加快训练我设置每10轮测试一次分割精度最后40轮每轮测试分割精度。代码修改指南梯度积累学习率和检测分割损失比后者在train.py内未暴露出来是一组非常重要的参数。必须清楚YOLOv5使用了梯度积累不管你的batchsize是多少“名义batchsize”都是作者预设的64。这意味着当你设置batchsize为16时将会每4个batch才更新一次参数具体积累几次看训练时候我打印的accumulate第一次是目标值后续是当前值即实际的batchsize是你输入的batchsize的倍数中最接近64的那个此处我将原代码修改为严格不超过64。因此你输入的batchsize17实际51要比16实际64小很多调整学习率时候应注意到这点。代码结构common.py此代码是YOLOv5中的常用基础操作类我在其中增加了BiSeNet的ARM、FFMRFB1、2非RFBNet魔改版本见代码注释ASPP接口上增加了用于砍通道的参数ASPPs先用1×1减少输入通道从而可以少砍一些中间通道Attention通道注意力相当于不带3×3卷积的ARM基础SEDAPPM见哈工大论文这里效果不明显PyramidPoolingPSPNet。yolo.pyYOLOv5的模型主架构代码包括Model类和检测要用的Detect类我把新加的四种分割头类放在这个代码中可能放common.py会更清爽些。所有新加的模块要放入模型都必须要经过Model类。模型配置文件models/yolov5s_city_seg.yaml模型配置文件可以看到我在检测层前面加了分割层配置并增加了分割类别Cityscapes是19。推理不同Head预训练模型不用修改想训练不同Head模型需要注释和解注释PSP、Base和Lab不用再改train.py但BiSe还要注释和解注释train.py的两个地方加入aux loss后续会说明接口设计缺陷但暂时没时间改实际上用PSP、Base、Lab就够了除非你想增加辅助损失。sml模型参照原版区别仅在控制深度和宽度的depth_multiplewidth_multiple数值BasePSP和Lab的分割头也会随sml自动放缩。自问自答训练时候报only test a child process或者加载卡住是BUG但是程序其实没有死等待一会若有谁解决了加载效率和这个问题烦请issue戳一下。testval mode加载器比较慢但正常不应该达到2分钟以上。没看过YOLO源码可以魔改吗- 可以直接在我的分割头和配置文件上改可以不用去看解析函数等代码。涉及修改backbone分割层不在24层接口不同的分割头时候需要至少读过yolo.py相关代码特别是Model的初始化、parse和run_once三部分。训这个前建议至少训过原版YOLOv5很多流程和原版一样。我把train_custom.py里的分割验证loader的batchsize改成1了目的是兼顾有不同尺寸的数据集同尺寸数据集嫌慢的可以手动改一下代码把分割验证batchsize调大。不同尺寸也嫌慢的可以去train_custom.py解注释val mode的loader然后把testval mode的loader注释。训自己数据集看上面的教程链接但还是建议读一下SegmentationDataset.py做了什么必要时候自己实现类。