YOLOV6之数据加载

一、数据加载

  位于engine.py文件下

self.train_loader, self.val_loader = self.get_data_loader(args, cfg, self.data_dict)

二、create_dataloader

  位于engine.py文件下的get_data_loader()函数中

train_loader = create_dataloader(train_path, args.img_size, 
                                 args.batch_size // args.world_size, grid_size,
                                 hyp=dict(cfg.data_aug), augment=True, rect=False, 
                                 rank=args.local_rank,
                                 workers=args.workers, shuffle=True, 
                                 check_images=args.check_images,
                                 check_labels=args.check_labels, 
                                 data_dict=data_dict, task='train')[0]

三、代码测试

from yolov6.data.data_load import create_dataloader
from yolov6.utils.config import Config
from yolov6.utils.events import load_yaml


if __name__ == "__main__":
    train_path = 'G:/coco/images/train2017'
    img_size = 640
    batch_size = 16
    grid_size = 32
    cfg = Config.fromfile('configs/yolov6s.py')  # get config
    data_dict = load_yaml('data/coco.yaml')      # load_yaml
    print(cfg.data_aug)
    print(data_dict)
    train_loader = create_dataloader(train_path, img_size, batch_size, grid_size,
                                     hyp=dict(cfg.data_aug), augment=True, rect=False, rank=-1,
                                     workers=4, shuffle=True, check_images=False,
                                     check_labels=False, data_dict=data_dict, task='train')[0]
    print(len(train_loader))

四、coco.yaml配置说明

  images和labels数据放置路径如下形式：

  配置参数如下：