专栏名称: 计算机视觉研究院
主要由来自于大学的研究生组成的团队,本平台从事机器学习与深度学习领域,主要在人脸检测与识别,多目标检测研究方向。本团队想通过计算机视觉战队平台打造属于自己的品牌,让更多相关领域的人了解本团队,结识更多相关领域的朋友,一起来学习,共同进步!
51好读  ›  专栏  ›  计算机视觉研究院

PyTorch版YOLOv4更新了,不仅适用于自定义数据集,还集成了注意力和MobileNet

计算机视觉研究院  · 公众号  ·  · 2020-10-10 07:00

正文

计算机视觉研究院专栏

作者:Edison_G

计算机视觉研究院
长按扫描二维码
关注我们 获取更多资讯

距离YOLOV4的推出 ,已经过去5个多月。YOLO 框架采用C语言作为底层代码,这对于惯用Python的研究者来说,实在是有点不友好。因此网上出现了很多基于各种深度学习框架的YOLO复现版本。近日,就有研究者在GitHub上更新了基于PyTorch的YOLOv4。

从今年4月YOLOv4发布后,对于这个目标检测框架,问的最多的问题或许就是:「有没有同学复现YOLOv4的, 可以交流一下么」。由于原版YOLO使用C语言进行编程,光凭这一点就让不少同学望而却步。网上有很多基于TF/Keras和Caffe等的复现版本,但不少项目只给了代码,并没有给出模型在COCO、PASCAL VOC数据集上的训练结果。
近日,有研究者在GitHub 上开源了一个项目:基于PyTorch深度学习框架的YOLOv4复现版本,该版本基于YOLOv4作者给出的实现AlexeyAB/darknet,并在PASCAL VOC、COCO和自定义数据集上运行。
项目地址: https://github.com/argusswift/YOLOv4-PyTorch
除此以外,该项目还向主干网络添加了一些有用的注意力方法,并实现了mobilenetv2-YOLOV4和mobilenetv3-YOLOV4。
Attentive YOLOv4
该项目向主干网络添加了一些注意力方法,如 SEnet、CBAM。

SEnet (CVPR 2017)

CBAM (CVPR 2018)
Mobilenet YOLOv4
该研究还实现了Mobilenetv2-YOLOV4 和Mobilenetv3-YOLOV4(只需更改config/yolov4_config.py中的MODEL_TYPE即可)。
下表展示了Mobilenetv2-YOLOV4的性能结果:

现在我们来看该项目的详细内容和要求。
环境要求
  • Nvida GeForce RTX 2080TI

  • CUDA10.0

  • CUDNN7.0

  • windows 或 linux 系统

  • python 3.6

特性
  • DO-Conv ( https://arxiv.org/abs/2006.12030 ) (torch>=1.2)

  • Attention

  • fp_16 training

  • Mish

  • Custom data

  • Data Augment (RandomHorizontalFlip, RandomCrop, RandomAffine, Resize)

  • Multi-scale Training (320 to 640)

  • focal loss

  • CIOU

  • Label smooth

  • Mixup

  • cosine lr

安装依赖项
运行脚本安装依赖项。你需要提供conda安装路径(例如~/anaconda3)以及所创建conda环境的名称(此处为YOLOv4-PyTorch)。
需要注意的是:安装脚本已在 Ubuntu 18.04 和 Window 10 系统上进行过测试。如果出现问题,请查看详细的安装说明:
https://github.com/argusswift/YOLOv4-PyTorch/blob/master/INSTALL.md
准备工作
1. git 复制 YOLOv4 库
准备工作的第一步是复制YOLOv4。
git clone github.com/argusswift/YOLOv4-PyTorch.git
然后更新配置文件「config/yolov4_config.py」中「PROJECT_PATH」。
2. 数据集准备
该项目准备了Pascal VOC和MSCOCO 2017数据集。其中PascalVOC 数据集包括VOC 2012_trainval、VOC 2007_trainval和VOC2007_test,MSCOCO 2017数据集包括train2017_img、train2017_ann、val2017_img、val2017_ann、test2017_img、test2017_list。
PascalVOC数据集下载命令:
# Download the data.cd $HOME/data




    
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tarwget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tarwget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar# Extract the data.tar -xvf VOCtrainval_11-May-2012.tartar -xvf VOCtrainval_06-Nov-2007.tartar -xvf VOCtest_06-Nov-2007.tar

MSCOCO 2017 数据集下载命令:

#step1: download the following data and annotation   2017 Train images [118K/18GB]   2017 Val images [5K/1GB]   2017 Test images [41K/6GB]   2017 Train/Val annotations [241MB]   #step2: arrange the data to the following structure   COCO   ---train   ---test   ---val   ---annotations
在数据集下载好后,需要进行以下操作:
  • 将数据集放入目录,更新config/yolov4_config.py中的DATA_PATH参数。

  • (对于COCO数据集)使用coco_to_voc.py将COCO数据类型转换为VOC数据类型。

  • 转换数据格式:使用utils/voc.py或utils/coco.py将 ]pascal voc*.xml格式(或COCO*.json格式)转换为*.txt格式(Image_path xmin0,ymin0,xmax0,ymax0,class0 xmin1,ymin1,xmax1,ymax1,class1 ...)。

3. 下载权重文件
1)darknet 预训练权重: yolov4
https://drive.google.com/file/d/1cewMfusmPjYWbrnuJRuKhPMwRe_b9PaT/view
2)Mobilenet 预训练权重:
mobilenetv2:
https://pan.baidu.com/share/init?surl=sjixK2L9L0YgQnvfDuVTJQ
提取码:args;
mobilenetv3:
https://pan.baidu.com/share/init?surl=75wKejULuM0ZD05b9iSftg
提取码:args。
3)在根目录下创建weight文件夹,将下载好的权重文件放到weight/目录下。
4)训练时在config/yolov4_config.py中设置MODEL_TYPE。
4. 转换成自定义数据集(基于自定义数据集进行训练)

1)将自定义数据集的图片放入JPEGImages文件夹,将注释文件放入Annotations文件夹。
2)使用xml_to_txt.py文件将训练和测试文件列表写入ImageSets/Main/*.txt。
3)转换数据格式:使用utils/voc.py或utils/coco.py将pascal voc*.xml格式(或COCO*.json格式)转换为*.txt格式(Image_path xmin0,ymin0,xmax0,ymax0,class0 xmin1,ymin1,xmax1,ymax1,class1 ...)。
训练
运行以下命令开始训练,详情参见config / yolov4_config.py。训练时应将DATA_TYPE设置为VOC或COCO。
CUDA_VISIBLE_DEVICES=0 nohup python -u train.py  --weight_path weight/yolov4.weights --gpu_id 0 > nohup.log 2>&1 &
它还支持resume训练,添加--resume,使用以下命令即可自动加载last.pt。
CUDA_VISIBLE_DEVICES=0 nohup python -u train.py  --weight_path weight/last.pt --gpu_id 0 > nohup.log 2>&1 &
检测
修改检测图像路径:DATA_TEST=/path/to/your/test_data# your own images。
for VOC dataset:CUDA_VISIBLE_DEVICES=0 python3 eval_voc.py --weight_path weight/best.pt --gpu_id 0 --visiual $DATA_TEST --eval --mode detfor COCO dataset:CUDA_VISIBLE_DEVICES=0 python3 eval_coco.py --weight_path weight/best.pt --gpu_id 0 --visiual $DATA_TEST --eval --mode det

结果可以在output/中查看,如下所示:

评估(Pascal VOC 数据集)
修改评估数据集路径:DATA_PATH=/path/to/your/test_data # your own images
for VOC dataset:CUDA_VISIBLE_DEVICES=0 python3 eval_voc.py --weight_path weight/best.pt --gpu_id 0 --visiual $DATA_TEST --eval -


评估(COCO 数据集)






请到「今天看啥」查看全文