(config)
| 56 | |
| 57 | |
| 58 | def build_loader(config): |
| 59 | config.defrost() |
| 60 | dataset_train, config.MODEL.NUM_CLASSES = build_dataset('train', config=config) |
| 61 | config.freeze() |
| 62 | print(f'local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}' |
| 63 | 'successfully build train dataset') |
| 64 | |
| 65 | dataset_val, _ = build_dataset('val', config=config) |
| 66 | print(f'local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}' |
| 67 | 'successfully build val dataset') |
| 68 | |
| 69 | dataset_test, _ = build_dataset('test', config=config) |
| 70 | print(f'local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}' |
| 71 | 'successfully build test dataset') |
| 72 | |
| 73 | num_tasks = dist.get_world_size() |
| 74 | global_rank = dist.get_rank() |
| 75 | |
| 76 | if dataset_train is not None: |
| 77 | if config.DATA.IMG_ON_MEMORY: |
| 78 | sampler_train = NodeDistributedSampler(dataset_train) |
| 79 | else: |
| 80 | if config.DATA.ZIP_MODE and config.DATA.CACHE_MODE == 'part': |
| 81 | indices = np.arange(dist.get_rank(), len(dataset_train), dist.get_world_size()) |
| 82 | sampler_train = SubsetRandomSampler(indices) |
| 83 | else: |
| 84 | sampler_train = torch.utils.data.DistributedSampler( |
| 85 | dataset_train, |
| 86 | num_replicas=num_tasks, |
| 87 | rank=global_rank, |
| 88 | shuffle=True) |
| 89 | |
| 90 | if dataset_val is not None: |
| 91 | if config.TEST.SEQUENTIAL: |
| 92 | sampler_val = torch.utils.data.SequentialSampler(dataset_val) |
| 93 | else: |
| 94 | sampler_val = torch.utils.data.distributed.DistributedSampler(dataset_val, shuffle=False) |
| 95 | |
| 96 | if dataset_test is not None: |
| 97 | if config.TEST.SEQUENTIAL: |
| 98 | sampler_test = torch.utils.data.SequentialSampler(dataset_test) |
| 99 | else: |
| 100 | sampler_test = torch.utils.data.distributed.DistributedSampler(dataset_test, shuffle=False) |
| 101 | |
| 102 | data_loader_train = torch.utils.data.DataLoader( |
| 103 | dataset_train, |
| 104 | sampler=sampler_train, |
| 105 | batch_size=config.DATA.BATCH_SIZE, |
| 106 | num_workers=config.DATA.NUM_WORKERS, |
| 107 | pin_memory=config.DATA.PIN_MEMORY, |
| 108 | drop_last=True, |
| 109 | persistent_workers=True) if dataset_train is not None else None |
| 110 | |
| 111 | data_loader_val = torch.utils.data.DataLoader( |
| 112 | dataset_val, |
| 113 | sampler=sampler_val, |
| 114 | batch_size=config.DATA.BATCH_SIZE, |
| 115 | shuffle=False, |
no test coverage detected