mirror of https://github.com/kritiksoman/GIMP-ML
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
32 lines
886 B
Python
32 lines
886 B
Python
import torch.utils.data
|
|
from data.base_data_loader import BaseDataLoader
|
|
|
|
|
|
def CreateDataset(opt):
|
|
dataset = None
|
|
from data.aligned_dataset import AlignedDataset
|
|
dataset = AlignedDataset()
|
|
|
|
print("dataset [%s] was created" % (dataset.name()))
|
|
dataset.initialize(opt)
|
|
return dataset
|
|
|
|
class CustomDatasetDataLoader(BaseDataLoader):
|
|
def name(self):
|
|
return 'CustomDatasetDataLoader'
|
|
|
|
def initialize(self, opt):
|
|
BaseDataLoader.initialize(self, opt)
|
|
self.dataset = CreateDataset(opt)
|
|
self.dataloader = torch.utils.data.DataLoader(
|
|
self.dataset,
|
|
batch_size=opt.batchSize,
|
|
shuffle=not opt.serial_batches,
|
|
num_workers=int(opt.nThreads))
|
|
|
|
def load_data(self):
|
|
return self.dataloader
|
|
|
|
def __len__(self):
|
|
return min(len(self.dataset), self.opt.max_dataset_size)
|