| from ..datasets.simmim_modis_dataset import MODISDataset | |
| from ..transforms import SimmimTransform | |
| import torch.distributed as dist | |
| from torch.utils.data import DataLoader, DistributedSampler | |
| from torch.utils.data._utils.collate import default_collate | |
| DATASETS = { | |
| 'MODIS': MODISDataset, | |
| } | |
| def collate_fn(batch): | |
| if not isinstance(batch[0][0], tuple): | |
| return default_collate(batch) | |
| else: | |
| batch_num = len(batch) | |
| ret = [] | |
| for item_idx in range(len(batch[0][0])): | |
| if batch[0][0][item_idx] is None: | |
| ret.append(None) | |
| else: | |
| ret.append(default_collate( | |
| [batch[i][0][item_idx] for i in range(batch_num)])) | |
| ret.append(default_collate([batch[i][1] for i in range(batch_num)])) | |
| return ret | |
| def get_dataset_from_dict(dataset_name): | |
| try: | |
| dataset_to_use = DATASETS[dataset_name] | |
| except KeyError: | |
| error_msg = f"{dataset_name} is not an existing dataset" | |
| error_msg = f"{error_msg}. Available datasets: {DATASETS.keys()}" | |
| raise KeyError(error_msg) | |
| return dataset_to_use | |
| def build_mim_dataloader(config, logger): | |
| transform = SimmimTransform(config) | |
| logger.info(f'Pre-train data transform:\n{transform}') | |
| dataset_name = config.DATA.DATASET | |
| dataset_to_use = get_dataset_from_dict(dataset_name) | |
| dataset = dataset_to_use(config, | |
| config.DATA.DATA_PATHS, | |
| split="train", | |
| img_size=config.DATA.IMG_SIZE, | |
| transform=transform) | |
| logger.info(f'Build dataset: train images = {len(dataset)}') | |
| sampler = DistributedSampler( | |
| dataset, | |
| num_replicas=dist.get_world_size(), | |
| rank=dist.get_rank(), | |
| shuffle=True) | |
| dataloader = DataLoader(dataset, | |
| config.DATA.BATCH_SIZE, | |
| sampler=sampler, | |
| num_workers=config.DATA.NUM_WORKERS, | |
| pin_memory=True, | |
| drop_last=True, | |
| collate_fn=collate_fn) | |
| return dataloader | |