123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100 |
- import importlib
- import numpy as np
- import random
- import torch
- import torch.utils.data
- from copy import deepcopy
- from functools import partial
- from os import path as osp
- from basicsr.data.prefetch_dataloader import PrefetchDataLoader
- from basicsr.utils import get_root_logger, scandir
- from basicsr.utils.dist_util import get_dist_info
- from basicsr.utils.registry import DATASET_REGISTRY
- __all__ = ['build_dataset', 'build_dataloader']
- # automatically scan and import dataset modules for registry
- # scan all the files under the data folder with '_dataset' in file names
- data_folder = osp.dirname(osp.abspath(__file__))
- dataset_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(data_folder) if v.endswith('_dataset.py')]
- # import all the dataset modules
- _dataset_modules = [importlib.import_module(f'basicsr.data.{file_name}') for file_name in dataset_filenames]
- def build_dataset(dataset_opt):
- """Build dataset from options.
- Args:
- dataset_opt (dict): Configuration for dataset. It must constain:
- name (str): Dataset name.
- type (str): Dataset type.
- """
- dataset_opt = deepcopy(dataset_opt)
- dataset = DATASET_REGISTRY.get(dataset_opt['type'])(dataset_opt)
- logger = get_root_logger()
- logger.info(f'Dataset [{dataset.__class__.__name__}] - {dataset_opt["name"]} ' 'is built.')
- return dataset
- def build_dataloader(dataset, dataset_opt, num_gpu=1, dist=False, sampler=None, seed=None):
- """Build dataloader.
- Args:
- dataset (torch.utils.data.Dataset): Dataset.
- dataset_opt (dict): Dataset options. It contains the following keys:
- phase (str): 'train' or 'val'.
- num_worker_per_gpu (int): Number of workers for each GPU.
- batch_size_per_gpu (int): Training batch size for each GPU.
- num_gpu (int): Number of GPUs. Used only in the train phase.
- Default: 1.
- dist (bool): Whether in distributed training. Used only in the train
- phase. Default: False.
- sampler (torch.utils.data.sampler): Data sampler. Default: None.
- seed (int | None): Seed. Default: None
- """
- phase = dataset_opt['phase']
- rank, _ = get_dist_info()
- if phase == 'train':
- if dist: # distributed training
- batch_size = dataset_opt['batch_size_per_gpu']
- num_workers = dataset_opt['num_worker_per_gpu']
- else: # non-distributed training
- multiplier = 1 if num_gpu == 0 else num_gpu
- batch_size = dataset_opt['batch_size_per_gpu'] * multiplier
- num_workers = dataset_opt['num_worker_per_gpu'] * multiplier
- dataloader_args = dict(
- dataset=dataset,
- batch_size=batch_size,
- shuffle=False,
- num_workers=num_workers,
- sampler=sampler,
- drop_last=True)
- if sampler is None:
- dataloader_args['shuffle'] = True
- dataloader_args['worker_init_fn'] = partial(
- worker_init_fn, num_workers=num_workers, rank=rank, seed=seed) if seed is not None else None
- elif phase in ['val', 'test']: # validation
- dataloader_args = dict(dataset=dataset, batch_size=1, shuffle=False, num_workers=0)
- else:
- raise ValueError(f'Wrong dataset phase: {phase}. ' "Supported ones are 'train', 'val' and 'test'.")
- dataloader_args['pin_memory'] = dataset_opt.get('pin_memory', False)
- prefetch_mode = dataset_opt.get('prefetch_mode')
- if prefetch_mode == 'cpu': # CPUPrefetcher
- num_prefetch_queue = dataset_opt.get('num_prefetch_queue', 1)
- logger = get_root_logger()
- logger.info(f'Use {prefetch_mode} prefetch dataloader: ' f'num_prefetch_queue = {num_prefetch_queue}')
- return PrefetchDataLoader(num_prefetch_queue=num_prefetch_queue, **dataloader_args)
- else:
- # prefetch_mode=None: Normal dataloader
- # prefetch_mode='cuda': dataloader for CUDAPrefetcher
- return torch.utils.data.DataLoader(**dataloader_args)
- def worker_init_fn(worker_id, num_workers, rank, seed):
- # Set the worker seed to num_workers * rank + worker_id + seed
- worker_seed = num_workers * rank + worker_id + seed
- np.random.seed(worker_seed)
- random.seed(worker_seed)
|