MLBench/ReaLens

1"""This package includes all the modules related to data loading and preprocessing2 3 To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.4 You need to implement four functions:5    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).6    -- <__len__>:                       return the size of dataset.7    -- <__getitem__>:                   get a data point from data loader.8    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.9 10Now you can use the dataset class by specifying flag '--dataset_mode dummy'.11See our template dataset class 'template_dataset.py' for more details.12"""13 14import importlib15import torch.utils.data16from torch.utils.data.distributed import DistributedSampler17import torch.distributed as dist18import os19from data.base_dataset import BaseDataset20 21 22def find_dataset_using_name(dataset_name):23    """Import the module "data/[dataset_name]_dataset.py".24 25    In the file, the class called DatasetNameDataset() will26    be instantiated. It has to be a subclass of BaseDataset,27    and it is case-insensitive.28    """29    dataset_filename = "data." + dataset_name + "_dataset"30    datasetlib = importlib.import_module(dataset_filename)31 32    dataset = None33    target_dataset_name = dataset_name.replace("_", "") + "dataset"34    for name, cls in datasetlib.__dict__.items():35        if name.lower() == target_dataset_name.lower() and issubclass(cls, BaseDataset):36            dataset = cls37 38    if dataset is None:39        raise NotImplementedError("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name))40 41    return dataset42 43 44def get_option_setter(dataset_name):45    """Return the static method <modify_commandline_options> of the dataset class."""46    dataset_class = find_dataset_using_name(dataset_name)47    return dataset_class.modify_commandline_options48 49 50def create_dataset(opt):51    """Create a dataset given the option.52 53    This function wraps the class CustomDatasetDataLoader.54        This is the main interface between this package and 'train.py'/'test.py'55 56    Example:57        >>> from data import create_dataset58        >>> dataset = create_dataset(opt)59    """60    data_loader = CustomDatasetDataLoader(opt)61    dataset = data_loader.load_data()62    return dataset63 64 65class CustomDatasetDataLoader:66    """Wrapper class of Dataset class that performs multi-threaded data loading"""67 68    def __init__(self, opt):69        """Initialize this class70 71        Step 1: create a dataset instance given the name [dataset_mode]72        Step 2: create a multi-threaded data loader.73        """74        self.opt = opt75        dataset_class = find_dataset_using_name(opt.dataset_mode)76        self.dataset = dataset_class(opt)77        print("dataset [%s] was created" % type(self.dataset).__name__)78 79        # Use DistributedSampler for DDP training80        if "LOCAL_RANK" in os.environ:81            print(f'create DDP sampler on rank {int(os.environ["LOCAL_RANK"])}')82            self.sampler = DistributedSampler(self.dataset, shuffle=not opt.serial_batches)83            shuffle = False  # DistributedSampler handles shuffling84        else:85            self.sampler = None86            shuffle = not opt.serial_batches87 88        self.dataloader = torch.utils.data.DataLoader(self.dataset, batch_size=opt.batch_size, shuffle=shuffle, sampler=self.sampler, num_workers=int(opt.num_threads))89 90    def load_data(self):91        return self92 93    def __len__(self):94        """Return the number of data in the dataset"""95        return min(len(self.dataset), self.opt.max_dataset_size)96 97    def __iter__(self):98        """Return a batch of data"""99        for i, data in enumerate(self.dataloader):100            if i * self.opt.batch_size >= self.opt.max_dataset_size:101                break102            yield data103 104    def set_epoch(self, epoch):105        """Set epoch for DistributedSampler to ensure proper shuffling"""106        if self.sampler is not None:107            self.sampler.set_epoch(epoch)108