2019-08-23 13:42:47 +00:00
|
|
|
"""create dataset and dataloader"""
|
|
|
|
import logging
|
|
|
|
import torch
|
|
|
|
import torch.utils.data
|
|
|
|
|
|
|
|
|
|
|
|
def create_dataloader(dataset, dataset_opt, opt=None, sampler=None):
|
|
|
|
phase = dataset_opt['phase']
|
|
|
|
if phase == 'train':
|
|
|
|
if opt['dist']:
|
|
|
|
world_size = torch.distributed.get_world_size()
|
|
|
|
num_workers = dataset_opt['n_workers']
|
|
|
|
assert dataset_opt['batch_size'] % world_size == 0
|
|
|
|
batch_size = dataset_opt['batch_size'] // world_size
|
|
|
|
shuffle = False
|
|
|
|
else:
|
2020-05-13 15:20:45 +00:00
|
|
|
num_workers = max(dataset_opt['n_workers'] * len(opt['gpu_ids']), 10)
|
2019-08-23 13:42:47 +00:00
|
|
|
batch_size = dataset_opt['batch_size']
|
|
|
|
shuffle = True
|
|
|
|
return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
|
|
|
|
num_workers=num_workers, sampler=sampler, drop_last=True,
|
|
|
|
pin_memory=False)
|
|
|
|
else:
|
2020-05-04 14:48:25 +00:00
|
|
|
batch_size = dataset_opt['batch_size'] or 1
|
2020-05-10 16:49:37 +00:00
|
|
|
return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=int(batch_size/2),
|
2019-08-23 13:42:47 +00:00
|
|
|
pin_memory=False)
|
|
|
|
|
|
|
|
|
|
|
|
def create_dataset(dataset_opt):
|
|
|
|
mode = dataset_opt['mode']
|
|
|
|
# datasets for image restoration
|
|
|
|
if mode == 'LQ':
|
|
|
|
from data.LQ_dataset import LQDataset as D
|
|
|
|
elif mode == 'LQGT':
|
|
|
|
from data.LQGT_dataset import LQGTDataset as D
|
2020-04-24 05:59:09 +00:00
|
|
|
# datasets for image corruption
|
|
|
|
elif mode == 'downsample':
|
|
|
|
from data.Downsample_dataset import DownsampleDataset as D
|
2019-08-23 13:42:47 +00:00
|
|
|
# datasets for video restoration
|
|
|
|
elif mode == 'REDS':
|
|
|
|
from data.REDS_dataset import REDSDataset as D
|
|
|
|
elif mode == 'Vimeo90K':
|
|
|
|
from data.Vimeo90K_dataset import Vimeo90KDataset as D
|
|
|
|
elif mode == 'video_test':
|
|
|
|
from data.video_test_dataset import VideoTestDataset as D
|
|
|
|
else:
|
|
|
|
raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode))
|
|
|
|
dataset = D(dataset_opt)
|
|
|
|
|
|
|
|
logger = logging.getLogger('base')
|
|
|
|
logger.info('Dataset [{:s} - {:s}] is created.'.format(dataset.__class__.__name__,
|
|
|
|
dataset_opt['name']))
|
|
|
|
return dataset
|