From 9054cf9b0c327a5077fd0793abe178f400da3315 Mon Sep 17 00:00:00 2001 From: tingchunw Date: Mon, 4 Dec 2017 16:52:46 -0800 Subject: first commit --- data/custom_dataset_data_loader.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100755 data/custom_dataset_data_loader.py (limited to 'data/custom_dataset_data_loader.py') diff --git a/data/custom_dataset_data_loader.py b/data/custom_dataset_data_loader.py new file mode 100755 index 0000000..0b98254 --- /dev/null +++ b/data/custom_dataset_data_loader.py @@ -0,0 +1,31 @@ +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) -- cgit v1.2.3-70-g09d2