| | |
| | import importlib |
| | import torch.utils.data |
| | from data.base_dataset import BaseDataset |
| |
|
| |
|
| | def find_dataset_using_name(dataset_name): |
| | dataset_filename = "data." + dataset_name + "_dataset" |
| |
|
| | datasetlib = importlib.import_module(dataset_filename) |
| |
|
| | dataset = None |
| | target_dataset_name = dataset_name.replace('_', '') + 'dataset' |
| | for name, cls in datasetlib.__dict__.items(): |
| | if name.lower() == target_dataset_name.lower() \ |
| | and issubclass(cls, BaseDataset): |
| | dataset = cls |
| |
|
| | if dataset is None: |
| | raise ValueError("In %s.py, there should be a subclass of BaseDataset " |
| | "with class name that matches %s in lowercase." % |
| | (dataset_filename, target_dataset_name)) |
| |
|
| | return dataset |
| |
|
| |
|
| | def get_option_setter(dataset_name): |
| | dataset_class = find_dataset_using_name(dataset_name) |
| | return dataset_class.modify_commandline_options |
| |
|
| |
|
| | def create_dataloader(opt): |
| | dataset = find_dataset_using_name(opt.dataset_mode) |
| | instance = dataset(opt) |
| |
|
| | print("dataset [%s] of size %d was created" % |
| | (type(instance).__name__, len(instance))) |
| | dataloader = torch.utils.data.DataLoader( |
| | instance, |
| | batch_size=opt.batchSize, |
| | shuffle=not opt.serial_batches, |
| | num_workers=int(opt.nThreads), |
| | drop_last=opt.isTrain |
| | ) |
| | return dataloader |
