文章目录
资源下载
(1) YOLOv6
Github
[美团技术团队] YOLOv6：又快又准的目标检测框架开源啦
(2) 参考博客
SIoU Loss: More Powerful Learning for Bounding Box Regression
YOLOX: Exceeding YOLO Series in 2021
RepVGG: Making VGG-style ConvNets Great Again
1 前言
本文主要记录使用 YOLOv6 训练自己数据集的过程，数据集以 Objects365 数据集为例.
2 数据集获取
链接：https://pan.baidu.com/s/1QiWm8hCJus3LstZkz6Mzdw
提取码：wmrx
3 数据集转化
Objects365 数据集为 COCO 格式数据，数据集文件格式如下:
1 | Objects365 |
YOLOv6 默认使用 YOLO 格式数据集，其中使用位置坐标格式为中心点坐标，数据集文件格式如下:
1 | Objects365_yolov6 |
通过以下脚本实现 COCO 格式的数据集转化为 YOLO 格式的数据集:
1 | #COCO 格式的数据集转化为 YOLO 格式的数据集 |
4 工程文件配置
(1) 配置模型文件
模型文件路径: config/yolov6n_objects365.py (新建) 以 YOLOv6n 为例:
1 | # YOLOv6n model |
(2) 配置数据集文件
数据集配置文件路径: data/objects365.yaml
1 | train: your/to/path/Objects365_yolov6/images/train2017 |
(3) 其他文件修改
针对 Objects365 此类类别个数超过超过 2 位数的数据集需修改以下文件:
文件路径: YOLOv6/yolov6/data/datasets.py
1 | @staticmethod |
5 模型训练及推理
(1) 训练环境搭建
使用 conda 或者 Docker 虚拟环境皆可，此处请自行搭建
(2) 模型训练
-
训练参数配置
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22# 文件位置: tools/train.py
def get_args_parser(add_help=True):
parser = argparse.ArgumentParser(description='YOLOv6 PyTorch Training', add_help=add_help)
parser.add_argument('--data-path', default='./data/coco.yaml', type=str, help='path of dataset')
parser.add_argument('--conf-file', default='./configs/yolov6s.py', type=str, help='experiments description file')
parser.add_argument('--img-size', type=int, default=640, help='train, val image size (pixels)')
parser.add_argument('--batch-size', default=32, type=int, help='total batch size for all GPUs')
parser.add_argument('--epochs', default=400, type=int, help='number of total epochs to run')
parser.add_argument('--workers', default=4, type=int, help='number of data loading workers (default: 8)')
parser.add_argument('--device', default='0', type=str, help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--eval-interval', type=int, default=20, help='evaluate at every interval epochs')
parser.add_argument('--eval-final-only', action='store_true', help='only evaluate at the final epoch')
parser.add_argument('--heavy-eval-range', default=50,
help='evaluating every epoch for last such epochs (can be jointly used with --eval-interval)')
parser.add_argument('--check-images', action='store_true', help='check images when initializing datasets')
parser.add_argument('--check-labels', action='store_true', help='check label files when initializing datasets')
parser.add_argument('--output-dir', default='./runs/train', type=str, help='path to save outputs')
parser.add_argument('--name', default='exp', type=str, help='experiment name, saved to output_dir/name')
parser.add_argument('--dist_url', type=str, default="default url: tcp://127.0.0.1:8888")
parser.add_argument('--gpu_count', type=int, default=0)
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter')
parser.add_argument('--resume', type=str, default=None, help='resume the corresponding ckpt') -
模型训练
1
python tools/train.py --batch-size 32 --conf-file configs/yolov6n_objects365.py --data-path data/objects365.yaml --device 0
(3) 模型验证
-
验证参数配置
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17# 文件位置: tools/eval.py
def get_args_parser(add_help=True):
parser = argparse.ArgumentParser(description='YOLOv6 PyTorch Evalating', add_help=add_help)
parser.add_argument('--data', type=str, default='./data/coco.yaml', help='dataset.yaml path')
parser.add_argument('--weights', type=str, default='./weights/yolov6s.pt', help='model.pt path(s)')
parser.add_argument('--batch-size', type=int, default=32, help='batch size')
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.65, help='NMS IoU threshold')
parser.add_argument('--task', default='val', help='val, or speed')
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--half', default=False, action='store_true', help='whether to use fp16 infer')
parser.add_argument('--save_dir', type=str, default='runs/val/', help='evaluation save dir')
parser.add_argument('--name', type=str, default='exp', help='save evaluation results to save_dir/name')
args = parser.parse_args()
LOGGER.info(args)
return args -
模型验证
1
python tools/eval.py --data data/objects365.yaml --batch-size 32 --weights weights/yolov6n.pt --task val
(4) 模型推理
-
推理参数配置
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24# 文件位置: tools/infer.py
def get_args_parser(add_help=True):
parser = argparse.ArgumentParser(description='YOLOv6 PyTorch Inference.', add_help=add_help)
parser.add_argument('--weights', type=str, default='weights/yolov6s.pt', help='model path(s) for inference.')
parser.add_argument('--source', type=str, default='data/images', help='the source path, e.g. image-file/dir.')
parser.add_argument('--yaml', type=str, default='data/coco.yaml', help='data yaml file.')
parser.add_argument('--img-size', type=int, default=640, help='the image-size(h,w) in inference size.')
parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold for inference.')
parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold for inference.')
parser.add_argument('--max-det', type=int, default=1000, help='maximal inferences per image.')
parser.add_argument('--device', default='0', help='device to run our model i.e. 0 or 0,1,2,3 or cpu.')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt.')
parser.add_argument('--save-img', action='store_false', help='save visuallized inference results.')
parser.add_argument('--classes', nargs='+', type=int, help='filter by classes, e.g. --classes 0, or --classes 0 2 3.')
parser.adkd_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS.')
parser.add_argument('--project', default='runs/inference', help='save inference results to project/name.')
parser.add_argument('--name', default='exp', help='save inference results to project/name.')
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels.')
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences.')
parser.add_argument('--half', action='store_true', help='whether to use FP16 half-precision inference.')
args = parser.parse_args()
LOGGER.info(args)
return args -
模型推理
1
python tools/infer.py --weights weights/yolov6n.pt --yaml data/objects365.yaml --source your/to/images.jpg
