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 (Dataset): A PyTorch dataset. samples_per_gpu (int): Number of training samples on each GPU, i.e.,
(dataset,
samples_per_gpu,
workers_per_gpu,
num_gpus=1,
dist=True,
shuffle=True,
seed=None,
shuffler_sampler=None,
nonshuffler_sampler=None,
**kwargs)
| 17 | from projects.mmdet3d_plugin.datasets.samplers.sampler import build_sampler |
| 18 | |
| 19 | def build_dataloader(dataset, |
| 20 | samples_per_gpu, |
| 21 | workers_per_gpu, |
| 22 | num_gpus=1, |
| 23 | dist=True, |
| 24 | shuffle=True, |
| 25 | seed=None, |
| 26 | shuffler_sampler=None, |
| 27 | nonshuffler_sampler=None, |
| 28 | **kwargs): |
| 29 | """Build PyTorch DataLoader. |
| 30 | In distributed training, each GPU/process has a dataloader. |
| 31 | In non-distributed training, there is only one dataloader for all GPUs. |
| 32 | Args: |
| 33 | dataset (Dataset): A PyTorch dataset. |
| 34 | samples_per_gpu (int): Number of training samples on each GPU, i.e., |
| 35 | batch size of each GPU. |
| 36 | workers_per_gpu (int): How many subprocesses to use for data loading |
| 37 | for each GPU. |
| 38 | num_gpus (int): Number of GPUs. Only used in non-distributed training. |
| 39 | dist (bool): Distributed training/test or not. Default: True. |
| 40 | shuffle (bool): Whether to shuffle the data at every epoch. |
| 41 | Default: True. |
| 42 | kwargs: any keyword argument to be used to initialize DataLoader |
| 43 | Returns: |
| 44 | DataLoader: A PyTorch dataloader. |
| 45 | """ |
| 46 | rank, world_size = get_dist_info() |
| 47 | if dist: |
| 48 | # DistributedGroupSampler will definitely shuffle the data to satisfy |
| 49 | # that images on each GPU are in the same group |
| 50 | if shuffle: |
| 51 | sampler = build_sampler(shuffler_sampler if shuffler_sampler is not None else dict(type='DistributedGroupSampler'), |
| 52 | dict( |
| 53 | dataset=dataset, |
| 54 | samples_per_gpu=samples_per_gpu, |
| 55 | num_replicas=world_size, |
| 56 | rank=rank, |
| 57 | seed=seed) |
| 58 | ) |
| 59 | |
| 60 | else: |
| 61 | sampler = build_sampler(nonshuffler_sampler if nonshuffler_sampler is not None else dict(type='DistributedSampler'), |
| 62 | dict( |
| 63 | dataset=dataset, |
| 64 | num_replicas=world_size, |
| 65 | rank=rank, |
| 66 | shuffle=shuffle, |
| 67 | seed=seed) |
| 68 | ) |
| 69 | |
| 70 | batch_size = samples_per_gpu |
| 71 | num_workers = workers_per_gpu |
| 72 | else: |
| 73 | # assert False, 'not support in bevformer' |
| 74 | print('WARNING!!!!, Only can be used for obtain inference speed!!!!') |
| 75 | sampler = GroupSampler(dataset, samples_per_gpu) if shuffle else None |
| 76 | batch_size = num_gpus * samples_per_gpu |