YOLOV5 训练

YOLOV5训练过程

CUDA 和cuDnnan 安装教程

windows上安装可以参考这篇知乎文章

数据集准备

自己准备数据集

可以使用 labelImg 工具，直接 pip install labelimg 就可以安装了。
命令行中输入 labelImg 就可以运行

标注数据的输出结果有多种过格式，VOC 、COCO 、YOLO等。

数据组织

先放目录树，建议先按照下面的目录格式，准备数据集。

└─VOCdevkit
    └─VOC2007
        ├─Annotations
        ├─JPEGImages

我们以VOC数据为例，说明数据集准备过程，我们需要将数据转成yolo格式。
VOC数据中，是以像素的绝对位置来标明 bndbox，但是yolo中的框是相对位置比例来标注的。
JPEGImages 文件夹中放图片，Annotions文件夹中放图片对应同名VOC数据（.xml 文件）。

下面的python代码，可以将数据集个数转为yolo数据格式，并自动划分train数据和val数据，划分比例通过 TRAIN_RATIO 调整。
使用之前首先修改 classes，将其修改为自己数据集中的类别。

注意该脚本应和VOCdevkit放在同一个目录中，就可以正常执行了。

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile
 
classes=["nest", "kite", "balloon", "trash"]

 
TRAIN_RATIO = 80
 
def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)
 
def convert(size, box):
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)
 
def convert_annotation(image_id):
    in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' %image_id)
    out_file = open('VOCdevkit/VOC2007/YOLOLabels/%s.txt' %image_id, 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    in_file.close()
    out_file.close()
wd = os.getcwd()
wd = os.getcwd()
data_base_dir = os.path.join(wd, "VOCdevkit/")
if not os.path.isdir(data_base_dir):
    os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
if not os.path.isdir(work_sapce_dir):
    os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
        os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
        os.mkdir(image_dir)
clear_hidden_files(image_dir)
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
if not os.path.isdir(yolo_labels_dir):
        os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)
yolov5_images_dir = os.path.join(data_base_dir, "images/")
if not os.path.isdir(yolov5_images_dir):
        os.mkdir(yolov5_images_dir)
clear_hidden_files(yolov5_images_dir)
yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
if not os.path.isdir(yolov5_labels_dir):
        os.mkdir(yolov5_labels_dir)
clear_hidden_files(yolov5_labels_dir)
yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
if not os.path.isdir(yolov5_images_train_dir):
        os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
if not os.path.isdir(yolov5_images_test_dir):
        os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
if not os.path.isdir(yolov5_labels_train_dir):
        os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)
yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
if not os.path.isdir(yolov5_labels_test_dir):
        os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)
 
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
train_file.close()
test_file.close()
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
list_imgs = os.listdir(image_dir) # list image files
prob = random.randint(1, 100)
print("Probability: %d" % prob)
for i in range(0,len(list_imgs)):
    path = os.path.join(image_dir,list_imgs[i])
    if os.path.isfile(path):
        image_path = image_dir + list_imgs[i]
        voc_path = list_imgs[i]
        (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
        (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
        annotation_name = nameWithoutExtention + '.xml'
        annotation_path = os.path.join(annotation_dir, annotation_name)
        label_name = nameWithoutExtention + '.txt'
        label_path = os.path.join(yolo_labels_dir, label_name)
    prob = random.randint(1, 100)
    print("Probability: %d" % prob)
    if(prob < TRAIN_RATIO): # train dataset
        if os.path.exists(annotation_path):
            train_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention) # convert label
            copyfile(image_path, yolov5_images_train_dir + voc_path)
            copyfile(label_path, yolov5_labels_train_dir + label_name)
    else: # test dataset
        if os.path.exists(annotation_path):
            test_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention) # convert label
            copyfile(image_path, yolov5_images_test_dir + voc_path)
            copyfile(label_path, yolov5_labels_test_dir + label_name)
train_file.close()
test_file.close()

执行完后的结果

└─VOCdevkit
    ├─images
    │  ├─train
    │  └─val
    ├─labels
    │  ├─train
    │  └─val
    └─VOC2007
        ├─Annotations
        ├─JPEGImages
        └─YOLOLabels

YOLOV5项目准备

下载YOLOV5源码 https://github.com/ultralytics/yolov5

├─.github
│  ├─ISSUE_TEMPLATE
│  └─workflows
├─.idea
│  └─inspectionProfiles
├─data
│  ├─images
│  └─scripts
├─models
│  ├─hub
│  └─__pycache__
├─utils
│  ├─aws
│  ├─google_app_engine
│  ├─wandb_logging
│  │  └─__pycache__
│  └─__pycache__
├─weights
└─__pycache__

在/data 文件夹中复制一个 .yaml 文件，重命名一下。将文件中的内容修改为自己的数据集。一般不将我们的训练数据直接放到 /data 文件夹中。

在/models 文件夹中复制一个 yolov5s.yaml 文件，重命名一下。将nc 修改一下，改成自己训练集中的类别数。

下载对应的预训练权重https://github.com/ultralytics/yolov5/releases, 下载对应的yolov5s.pt 就可以了

启动训练

python train.py --img 640 --batch 16 --epochs 300 --data ./data/myData.yaml --cfg ./models/myYolov5s.yaml --weights ./weights/yolov5s.pt  --device 0

训练完成后，会输出最好的模型参数存放位置。

训练过程报错

如果您在使用 PyTorch 1.7.1 或更高版本时遇到了 ModuleNotFoundError: No module named ‘torch.ao’ 的错误，可能是因为您的代码中使用了 torch.ao 模块，但该模块在您的环境中并不存在。

找到引用 torch.ao 的地方，将其改为 torch。

推理测试

python detect.py --source xxxxx.jpg --we ights ./weights/xxxx.pt --output inference/2_output/1_img/ --device 0

推理测试输出的结果没有框

1 在detec.py 中缺少一行 cudnn.benchmark=True

    if webcam:
        view_img = check_imshow()
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz, stride=stride)
    else:
        cudnn.benchmark = True
        dataset = LoadImages(source, img_size=imgsz, stride=stride)

2 可能是训练的准确度不够，需要调低阈值
减小下图中的值。