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

计算机视觉研究院

发布于 2020-10-19 10:09:16

5750

发布于 2020-10-19 10:09:16

作者：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.tar
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar
# Extract the data.
tar -xvf VOCtrainval_11-May-2012.tar
tar -xvf VOCtrainval_06-Nov-2007.tar
tar -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 det
for 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 数据集）

修改评估数据集路径：DATA_PATH=/path/to/your/test_data # your own images

CUDA_VISIBLE_DEVICES=0 python3 eval_coco.py --weight_path weight/best.pt --gpu_id 0 --visiual $DATA_TEST --eval --mode val

type=bbox
Running per image evaluation...      DONE (t=0.34s).
Accumulating evaluation results...   DONE (t=0.08s). 
Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.438  
Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.607  
Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.469  
Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.253  
Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.486  
Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.567  
Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.342  
Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.571  
Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.632  
Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.458  
Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.691  
Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.790

可视化热图

在 val_voc.py 中设置 showatt=Ture，网络即可输出热图。

for VOC dataset:
CUDA_VISIBLE_DEVICES=0 python3 eval_voc.py --weight_path weight/best.pt --gpu_id 0 --visiual $DATA_TEST --eval
for COCO dataset:
CUDA_VISIBLE_DEVICES=0 python3 eval_coco.py --weight_path weight/best.pt --gpu_id 0 --visiual $DATA_TEST --eval

在 output / 中可以查看热图，如下所示：