YOLOv4-pytorch训练自己的数据集

YOLOv4-pytorch介绍

Github地址：argusswift/YOLOv4-pytorch：https://github.com/argusswift/YOLOv4-pytorch
这是基于darknet的YOLOv4结构的PyTorch版的复现，还提供了Mobilenetv3-YOLOv4、attentive YOLOv4等有用的模块，操作简便，易读性强。

环境配置

运行环境

• Nvida GeForce RTX 2080TI
• CUDA10.0
• CUDNN7.0
• windows or linux
• python 3.6

安装依赖包

pip3 install -r requirements.txt --user

准备工作

Git clone YOLOv4

git clone github.com/argusswift/YOLOv4-pytorch.git

准备数据集

该模型提供了三种支持数据格式（PASCAL VOC, COCO, Customer）。

下载PascalVOC/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

1. 将数据集放至目录下，并更新config/yolov4_config.py中的"DATA_PATH"为数据集位置；
2. （对于COCO数据集）使用utils/coco_to_voc.py将COCO数据类型转化为VOC数据类型；
3. 使用utils/voc.py将PascalVOC的*.xml格式转化为*.txt格式或utils/coco.py将COCO的*.json格式转化为*.txt格式(Image_path xmin0,ymin0,xmax0,ymax0,class0 xmin1,ymin1,xmax1,ymax1,class1 …)。

准备自己的数据集

类似PascalVOC类型构建自己的数据集：

- VOC
	- JPEGImage #原图片文件
	- Annotations #标注*.xml文件
	- ImageSets
		 - Main #训练、测试集
			 - train.txt
			 - test.txt

1. 将图片放于JPEGImage文件夹下，标注文件放于Annotations文件夹下；
2. 使用utils/xml_to_txt.py将训练集和测试集写入ImageSets/Main/*.txt;
3. 使用utils/voc.py将PascalVOC的*.xml格式转化为*.txt格式或utils/coco.py将COCO的*.json格式转化为*.txt格式(Image_path xmin0,ymin0,xmax0,ymax0,class0 xmin1,ymin1,xmax1,ymax1,class1 …)；
4. 修改config/yolov4_config.py中Customer_DATA的NUM和CLASSES。

Customer_DATA = {
    
    
    "NUM": 2,  # your dataset number
    "CLASSES": [
        "name",
        "flag"
    ],  # your dataset class
}

下载权重文件

1. Darknet预训练权重：YOLOv4；
2. Mobilenet预训练权重：mobilenetv2, mobilenetv3（解压密码：args）；
3. 新建文件夹weight/，将权重文件放入；
4. 在config/yolov4_config.py中修改MODEL_TYPE。

MODEL_TYPE = {
    
    
    "TYPE": "YOLOv4"
}  # YOLO type:YOLOv4, Mobilenet-YOLOv4 or Mobilenetv3-YOLOv4

训练

修改config/yolov4_config.py中的参数：

TRAIN = {
    
    
    "DATA_TYPE": "Customer",  # DATA_TYPE: VOC ,COCO or Customer
    "TRAIN_IMG_SIZE": 416,
    "AUGMENT": True,
    "BATCH_SIZE": 8,
    "MULTI_SCALE_TRAIN": False,
    "IOU_THRESHOLD_LOSS": 0.5,
    "YOLO_EPOCHS": 4000,
    "Mobilenet_YOLO_EPOCHS": 120,
    "NUMBER_WORKERS": 0,
    "MOMENTUM": 0.9,
    "WEIGHT_DECAY": 0.0005,
    "LR_INIT": 1e-4,
    "LR_END": 1e-6,
    "WARMUP_EPOCHS": 2,  # or None
}

训练指令：

python -u train.py  --weight_path weight/yolov4.weights --gpu_id 0

或（nohup）

CUDA_VISIBLE_DEVICES=0 nohup python -u train.py  --weight_path weight/yolov4.weights --gpu_id 0 > nohup.log 2>&1 &

或（使用–resume，自动调用last.pt）

CUDA_VISIBLE_DEVICES=0 nohup python -u train.py  --weight_path weight/last.pt --gpu_id 0 > nohup.log 2>&1 &

测试

图片测试

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

视频测试

CUDA_VISIBLE_DEVICES=0 python3 video_test.py --weight_path best.pt --gpu_id 0 --video_path video.mp4 --output_dir --output_dir

遇到的问题

1. evaluater.py找不到*.xml标注文件

FileNotFoundError: [Errno 2] No such file or directory: '/home/my/YOLOv4-pytorch/data/VOC/Annotations\\18_3_dets0.xml'

报错原因：路径地址不正确
解决方法：
1.检查yolov4_config.py中DATA_PATH地址是否正确
2.evaluater.py，221 改为 self.val_data_path, "Annotations/" + "{:s}.xml"

参考文献：
[1]: https://github.com/argusswift/YOLOv4-pytorch