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)
| 54 | |
| 55 | |
| 56 | def build_dataloader(dataset: Dataset, |
| 57 | samples_per_gpu: int, |
| 58 | workers_per_gpu: int, |
| 59 | num_gpus: Optional[int] = 1, |
| 60 | dist: Optional[bool] = True, |
| 61 | shuffle: Optional[bool] = True, |
| 62 | round_up: Optional[bool] = True, |
| 63 | seed: Optional[Union[int, None]] = None, |
| 64 | persistent_workers: Optional[bool] = True, |
| 65 | **kwargs): |
| 66 | """Build PyTorch DataLoader. |
| 67 | |
| 68 | In distributed training, each GPU/process has a dataloader. |
| 69 | In non-distributed training, there is only one dataloader for all GPUs. |
| 70 | |
| 71 | Args: |
| 72 | dataset (:obj:`Dataset`): A PyTorch dataset. |
| 73 | samples_per_gpu (int): Number of training samples on each GPU, i.e., |
| 74 | batch size of each GPU. |
| 75 | workers_per_gpu (int): How many subprocesses to use for data loading |
| 76 | for each GPU. |
| 77 | num_gpus (int, optional): Number of GPUs. Only used in non-distributed |
| 78 | training. |
| 79 | dist (bool, optional): Distributed training/test or not. Default: True. |
| 80 | shuffle (bool, optional): Whether to shuffle the data at every epoch. |
| 81 | Default: True. |
| 82 | round_up (bool, optional): Whether to round up the length of dataset by |
| 83 | adding extra samples to make it evenly divisible. Default: True. |
| 84 | persistent_workers (bool): If True, the data loader will not shutdown |
| 85 | the worker processes after a dataset has been consumed once. |
| 86 | This allows to maintain the workers Dataset instances alive. |
| 87 | The argument also has effect in PyTorch>=1.7.0. |
| 88 | Default: True |
| 89 | kwargs: any keyword argument to be used to initialize DataLoader |
| 90 | |
| 91 | Returns: |
| 92 | DataLoader: A PyTorch dataloader. |
| 93 | """ |
| 94 | rank, world_size = get_dist_info() |
| 95 | if dist: |
| 96 | sampler = DistributedSampler( |
| 97 | dataset, world_size, rank, shuffle=shuffle, round_up=round_up) |
| 98 | shuffle = False |
| 99 | batch_size = samples_per_gpu |
| 100 | num_workers = workers_per_gpu |
| 101 | else: |
| 102 | sampler = None |
| 103 | batch_size = num_gpus * samples_per_gpu |
| 104 | num_workers = num_gpus * workers_per_gpu |
| 105 | |
| 106 | init_fn = partial( |
| 107 | worker_init_fn, num_workers=num_workers, rank=rank, |
| 108 | seed=seed) if seed is not None else None |
| 109 | |
| 110 | data_loader = DataLoader( |
| 111 | dataset, |
| 112 | batch_size=batch_size, |
| 113 | sampler=sampler, |
no test coverage detected