Build PyTorch DataLoader. In distributed training, each GPU/process has a dataloader. In non-distributed training, there is only one dataloader for all GPUs. Args: dataset (:obj:`Dataset`): A PyTorch dataset. samples_per_gpu (int): Number of training samples on each GPU,
(dataset: Dataset,
samples_per_gpu: int,
workers_per_gpu: int,
num_gpus: Optional[int] = 1,
dist: Optional[bool] = True,
shuffle: Optional[bool] = True,
round_up: Optional[bool] = True,
seed: Optional[Union[int, None]] = None,
persistent_workers: Optional[bool] = True,
**kwargs)
| 41 | |
| 42 | |
| 43 | def build_dataloader(dataset: Dataset, |
| 44 | samples_per_gpu: int, |
| 45 | workers_per_gpu: int, |
| 46 | num_gpus: Optional[int] = 1, |
| 47 | dist: Optional[bool] = True, |
| 48 | shuffle: Optional[bool] = True, |
| 49 | round_up: Optional[bool] = True, |
| 50 | seed: Optional[Union[int, None]] = None, |
| 51 | persistent_workers: Optional[bool] = True, |
| 52 | **kwargs): |
| 53 | """Build PyTorch DataLoader. |
| 54 | In distributed training, each GPU/process has a dataloader. |
| 55 | In non-distributed training, there is only one dataloader for all GPUs. |
| 56 | Args: |
| 57 | dataset (:obj:`Dataset`): A PyTorch dataset. |
| 58 | samples_per_gpu (int): Number of training samples on each GPU, i.e., |
| 59 | batch size of each GPU. |
| 60 | workers_per_gpu (int): How many subprocesses to use for data loading |
| 61 | for each GPU. |
| 62 | num_gpus (int, optional): Number of GPUs. Only used in non-distributed |
| 63 | training. |
| 64 | dist (bool, optional): Distributed training/test or not. Default: True. |
| 65 | shuffle (bool, optional): Whether to shuffle the data at every epoch. |
| 66 | Default: True. |
| 67 | round_up (bool, optional): Whether to round up the length of dataset by |
| 68 | adding extra samples to make it evenly divisible. Default: True. |
| 69 | kwargs: any keyword argument to be used to initialize DataLoader |
| 70 | Returns: |
| 71 | DataLoader: A PyTorch dataloader. |
| 72 | """ |
| 73 | rank, world_size = get_dist_info() |
| 74 | if dist: |
| 75 | sampler = DistributedSampler(dataset, |
| 76 | world_size, |
| 77 | rank, |
| 78 | shuffle=shuffle, |
| 79 | round_up=round_up) |
| 80 | shuffle = False |
| 81 | batch_size = samples_per_gpu |
| 82 | num_workers = workers_per_gpu |
| 83 | else: |
| 84 | sampler = None |
| 85 | batch_size = num_gpus * samples_per_gpu |
| 86 | num_workers = num_gpus * workers_per_gpu |
| 87 | |
| 88 | init_fn = partial( |
| 89 | worker_init_fn, num_workers=num_workers, rank=rank, |
| 90 | seed=seed) if seed is not None else None |
| 91 | |
| 92 | data_loader = DataLoader(dataset, |
| 93 | batch_size=batch_size, |
| 94 | sampler=sampler, |
| 95 | num_workers=num_workers, |
| 96 | collate_fn=partial( |
| 97 | collate, samples_per_gpu=samples_per_gpu), |
| 98 | pin_memory=False, |
| 99 | shuffle=shuffle, |
| 100 | worker_init_fn=init_fn, |