(config)
| 23 | |
| 24 | |
| 25 | def build_loader(config): |
| 26 | config.defrost() |
| 27 | dataset_train, config.MODEL.NUM_CLASSES = build_dataset(is_train=True, config=config) |
| 28 | config.freeze() |
| 29 | print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build train dataset") |
| 30 | dataset_val, _ = build_dataset(is_train=False, config=config) |
| 31 | print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build val dataset") |
| 32 | |
| 33 | num_tasks = dist.get_world_size() |
| 34 | global_rank = dist.get_rank() |
| 35 | if config.DATA.ZIP_MODE and config.DATA.CACHE_MODE == 'part': |
| 36 | indices = np.arange(dist.get_rank(), len(dataset_train), dist.get_world_size()) |
| 37 | sampler_train = SubsetRandomSampler(indices) |
| 38 | else: |
| 39 | sampler_train = torch.utils.data.DistributedSampler( |
| 40 | dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True |
| 41 | ) |
| 42 | |
| 43 | indices = np.arange(dist.get_rank(), len(dataset_val), dist.get_world_size()) |
| 44 | sampler_val = SubsetRandomSampler(indices) |
| 45 | |
| 46 | data_loader_train = torch.utils.data.DataLoader( |
| 47 | dataset_train, sampler=sampler_train, |
| 48 | batch_size=config.DATA.BATCH_SIZE, |
| 49 | num_workers=config.DATA.NUM_WORKERS, |
| 50 | pin_memory=config.DATA.PIN_MEMORY, |
| 51 | drop_last=True, |
| 52 | ) |
| 53 | |
| 54 | data_loader_val = torch.utils.data.DataLoader( |
| 55 | dataset_val, sampler=sampler_val, |
| 56 | batch_size=config.DATA.BATCH_SIZE, |
| 57 | shuffle=False, |
| 58 | num_workers=config.DATA.NUM_WORKERS, |
| 59 | pin_memory=config.DATA.PIN_MEMORY, |
| 60 | drop_last=False |
| 61 | ) |
| 62 | |
| 63 | # setup mixup / cutmix |
| 64 | mixup_fn = None |
| 65 | mixup_active = config.AUG.MIXUP > 0 or config.AUG.CUTMIX > 0. or config.AUG.CUTMIX_MINMAX is not None |
| 66 | if mixup_active: |
| 67 | mixup_fn = Mixup( |
| 68 | mixup_alpha=config.AUG.MIXUP, cutmix_alpha=config.AUG.CUTMIX, cutmix_minmax=config.AUG.CUTMIX_MINMAX, |
| 69 | prob=config.AUG.MIXUP_PROB, switch_prob=config.AUG.MIXUP_SWITCH_PROB, mode=config.AUG.MIXUP_MODE, |
| 70 | label_smoothing=config.MODEL.LABEL_SMOOTHING, num_classes=config.MODEL.NUM_CLASSES) |
| 71 | |
| 72 | return dataset_train, dataset_val, data_loader_train, data_loader_val, mixup_fn |
| 73 | |
| 74 | |
| 75 | def build_dataset(is_train, config): |
no test coverage detected