Add random_dataset for testing

This commit is contained in:
James Betker 2020-12-09 14:55:05 -07:00
parent 97ff25a086
commit 66cbae8731
2 changed files with 19 additions and 0 deletions

View File

@ -49,6 +49,8 @@ def create_dataset(dataset_opt):
from data.torch_dataset import TorchDataset as D
elif mode == 'byol_dataset':
from data.byol_attachment import ByolDatasetWrapper as D
elif mode == 'random_dataset':
from data.random_dataset import RandomDataset as D
else:
raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode))
dataset = D(dataset_opt)

View File

@ -0,0 +1,17 @@
import torch
from torch.utils.data import Dataset
# Dataset that feeds random data into the state. Can be useful for testing or demo purposes without actual data.
class RandomDataset(Dataset):
def __init__(self, opt):
self.hq_shape = tuple(opt['hq_shape'])
self.lq_shape = tuple(opt['lq_shape'])
def __getitem__(self, item):
return {'lq': torch.rand(self.lq_shape), 'hq': torch.rand(self.hq_shape),
'LQ_path': '', 'GT_path': ''}
def __len__(self):
# Arbitrary
return 1024 * 1024