YOLOv4-pytorch训练自己的数据集
YOLOv4-pytorch介绍
Github地址: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: 、、 MSCOCO 2017: 、 、 、 、 、
- 将数据集放至目录下,并更新config/yolov4_config.py中的"DATA_PATH"为数据集位置;
- (对于COCO数据集)使用utils/coco_to_voc.py将COCO数据类型转化为VOC数据类型;
- 使用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
- 将图片放于JPEGImage文件夹下,标注文件放于Annotations文件夹下;
- 使用utils/xml_to_txt.py将训练集和测试集写入ImageSets/Main/*.txt;
- 使用utils/voc.py将PascalVOC的*.xml格式转化为*.txt格式或utils/coco.py将COCO的*.json格式转化为*.txt格式(Image_path xmin0,ymin0,xmax0,ymax0,class0 xmin1,ymin1,xmax1,ymax1,class1 …);
- 修改config/yolov4_config.py中Customer_DATA的NUM和CLASSES。
Customer_DATA = { "NUM": 2, # your dataset number "CLASSES": [ "name", "flag" ], # your dataset class }
下载权重文件
- Darknet预训练权重:;
- Mobilenet预训练权重:, (解压密码:args);
- 新建文件夹weight/,将权重文件放入;
- 在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
遇到的问题
- 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]:
上一篇:
通过多线程提高代码的执行效率例子
下一篇:
Anaconda安装教程(超详细)