YOLOV6训练自己的数据集

开局一座山，算法数据全靠搬，开整默认你已经准备好数据data(图像)与data_labels（不知道怎么生成的可以看我之前的博客训练yolov5的txt)

开源地址：https://github.com/meituan/YOLOv6​​​​​​​

第一步：把冰箱门打开

准备数据根据下面的示例组织您的 train 和 val 图像和标签文件。

path/to/data/images/train/im0.jpg
path/to/data/images/val/im1.jpg
path/to/data/images/test/im2.jpg

# label directory
path/to/data/labels/train/im0.txt
path/to/data/labels/val/im1.txt
path/to/data/labels/test/im2.txt

这里我们采用脚本生成如下：

# conding:utf-8
import sys,os
import math
import random
import shutil
root_path = sys.path[0]
values_path = root_path+"/data"
values_list = os.listdir(values_path)

# 生成训练集、测试集、验证集，放文件的位置
target_path = '/home/project/data/images'
target_label_path = '/home/project/data/labels'
target_path_list = [target_path,target_label_path]
train_and_val = random.sample(values_list,int(math.floor(len(values_list)*0.9)))
train_list = random.sample(train_and_val,int(math.floor(len(train_and_val)*8/9)))
val_list = list(set(train_and_val).difference(set(train_list)))
test_list = list(set(values_list).difference(set(train_and_val)))
def create_set(input_list_name, folder_name):
    # 创建相应的文件夹
    for target_g_path in target_path_list:
        folder_name_path = target_g_path + "/" + folder_name
        if not os.path.exists(folder_name_path):
            os.mkdir(folder_name_path)
    # 将数据和标签数复制到指定的位置
    for file in eval(input_list_name):
        # 1.数据写入
        # 数据源所在文件夹
        source_path = root_path + "/data"
        full_target_file_name = os.path.join(source_path, file)  # 把文件的完整路径得到
        # 写入到指定的文件
        target_folder_path = os.path.join(target_path,folder_name)
        shutil.copy(full_target_file_name, target_folder_path)

        # 数据标签写入
        source_path = root_path + "/data_label"
        file = file.strip().split('.')[0] + '.txt'
        full_target_file_name = os.path.join(source_path, file)  # 把文件的完整路径得到
        target_folder_path = os.path.join(target_label_path, folder_name)
        shutil.copy(full_target_file_name, target_folder_path)



key_name = ['train', 'val', 'test']
for folder_name in key_name:
    list_name = folder_name + "_list"
    ##txt_name = folder_name + ".txt"
    create_set(list_name, folder_name)

第二步把东西塞进去：修改参数

 

dataset.yaml在$YOLOv6_DIR/data

train: path/to/data/images/train # train images
val: path/to/data/images/val # val images
test: path/to/data/images/test # test images (optional)

# Classes
nc: 1  # number of classes
names: ['apple']  # class names

创建配置文件

我们使用配置文件来指定网络结构和训练设置，包括优化器和数据增强超参数。

如果你创建一个新的配置文件，请将它放在 configs 目录下。或者只是使用提供的配置文件在$YOLOV6_HOME/configs/*_finetune.py

## YOLOv6s Model config file
model = dict(
    type='YOLOv6s',
    pretrained='./weights/yolov6s.pt', # download pretrain model from YOLOv6 github if use pretrained model
    depth_multiple = 0.33,
    width_multiple = 0.50,
    ...
)
solver=dict(
    optim='SGD',
    lr_scheduler='Cosine',
    ...
)

data_aug = dict(
    hsv_h=0.015,
    hsv_s=0.7,
    hsv_v=0.4,
    ...
)

第三步：关门

训练

单 GPU：默认的是400个周期

python tools/train.py --batch 32 --conf configs/yolov6s_finetune.py --data data/data.yaml --device 0

多 GPU（推荐 DDP 模式）

python -m torch.distributed.launch --nproc_per_node 4 tools/train.py --batch 256 --conf configs/yolov6s_finetune.py --data data/data.yaml --device 0,1,2,3

 评价

python tools/eval.py --data data/data.yaml  --weights output_dir/name/weights/best_ckpt.pt --device 0

推理

python tools/infer.py --weights output_dir/name/weights/best_ckpt.pt --source img.jpg --device 0

部署

导出为 ONNX 格式

python deploy/ONNX/export_onnx.py --weights output_dir/name/weights/best_ckpt.pt --device 0