PyTorch版YOLOv4更新了,适用于自定义数据集

开发 前端
距离YOLO v4 的推出,已经过去 5 个多月。YOLO 框架采用 C 语言作为底层代码,这对于惯用 Python 的研究者来说,实在是有点不友好。

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

[[344844]]

从今年 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。

PyTorch版YOLOv4更新了,适用于自定义数据集

SEnet (CVPR 2017)

PyTorch版YOLOv4更新了,适用于自定义数据集

CBAM (CVPR 2018)

mobilenet YOLOv4

该研究还实现了 mobilenetv2-YOLOV4 和 mobilenetv3-YOLOV4(只需更改 config/yolov4_config.py 中的 MODEL_TYPE 即可)。

下表展示了 mobilenetv2-YOLOV4 的性能结果:

PyTorch版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。

然后更新配置文件「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 数据集下载命令:

MSCOCO 2017 数据集下载命令:

在数据集下载好后,需要进行以下操作:

将数据集放入目录,更新 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。

它还支持 resume 训练,添加 --resume,使用以下命令即可自动加载 last.pt。

检测

修改检测图像路径:DATA_TEST=/path/to/your/test_data# your own images。

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

PyTorch版YOLOv4更新了,适用于自定义数据集

评估(Pascal VOC 数据集)

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

PyTorch版YOLOv4更新了,适用于自定义数据集

评估(COCO 数据集)

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

可视化热图

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

在 output / 中可以查看热图,如下所示:

PyTorch版YOLOv4更新了,适用于自定义数据集

 

责任编辑:张燕妮 来源: 机器之心Pro
相关推荐

2018-12-13 11:19:21

2020-06-12 14:25:36

框架PyTorch开发

2021-08-30 09:00:00

人工智能计算机视觉机器学习

2011-08-17 11:08:59

2021-05-10 15:00:32

Lumia 950Windows 10镜像

2021-05-05 20:50:56

苹果SafariSafari 14.1

2018-11-12 11:50:10

开源发票工具Web

2023-09-03 12:52:17

2011-08-01 16:10:11

XCode Excel 数据库

2015-08-25 15:02:51

UbuntuLinux发行版

2023-09-12 13:59:41

OpenAI数据集

2023-01-12 12:20:29

YOLOv8模型

2011-02-25 09:08:22

LinuxSuSE Linux

2022-09-02 17:47:46

Linux笔记应用

2021-11-04 10:42:43

汽车软件技术

2011-03-11 15:53:07

CentOS安装LAMP

2011-12-08 09:43:56

虚拟化vmwareVMware Fusi

2023-11-30 08:55:15

LinuxLibreOffic

2020-03-13 10:46:35

数据分析商业智能BI

2018-10-07 05:08:11

点赞
收藏

51CTO技术栈公众号