get_data_loader returns torch.utils.data.DataLoader for a Dataset. All arguments are comparable with those of pytorch DataLoader. However, if distributed, DistributedProxySampler, which is a wrapper of data_sampler, is used. Args num_epochs: total batch -> (# of batches
(dset,
batch_size=None,
shuffle=False,
num_workers=4,
pin_memory=False,
data_sampler=None,
replacement=True,
num_epochs=None,
num_iters=None,
generator=None,
drop_last=True,
distributed=False)
| 83 | |
| 84 | |
| 85 | def get_data_loader(dset, |
| 86 | batch_size=None, |
| 87 | shuffle=False, |
| 88 | num_workers=4, |
| 89 | pin_memory=False, |
| 90 | data_sampler=None, |
| 91 | replacement=True, |
| 92 | num_epochs=None, |
| 93 | num_iters=None, |
| 94 | generator=None, |
| 95 | drop_last=True, |
| 96 | distributed=False): |
| 97 | """ |
| 98 | get_data_loader returns torch.utils.data.DataLoader for a Dataset. |
| 99 | All arguments are comparable with those of pytorch DataLoader. |
| 100 | However, if distributed, DistributedProxySampler, which is a wrapper of data_sampler, is used. |
| 101 | |
| 102 | Args |
| 103 | num_epochs: total batch -> (# of batches in dset) * num_epochs |
| 104 | num_iters: total batch -> num_iters |
| 105 | """ |
| 106 | |
| 107 | assert batch_size is not None |
| 108 | |
| 109 | if data_sampler is None: |
| 110 | return DataLoader(dset, batch_size=batch_size, shuffle=shuffle, |
| 111 | num_workers=num_workers, pin_memory=pin_memory) |
| 112 | |
| 113 | else: |
| 114 | if isinstance(data_sampler, str): |
| 115 | data_sampler = get_sampler_by_name(data_sampler) |
| 116 | |
| 117 | if distributed: |
| 118 | assert dist.is_available() |
| 119 | num_replicas = dist.get_world_size() |
| 120 | else: |
| 121 | num_replicas = 1 |
| 122 | |
| 123 | if (num_epochs is not None) and (num_iters is None): |
| 124 | num_samples = len(dset) * num_epochs |
| 125 | elif (num_epochs is None) and (num_iters is not None): |
| 126 | num_samples = batch_size * num_iters * num_replicas |
| 127 | else: |
| 128 | num_samples = len(dset) |
| 129 | |
| 130 | if data_sampler.__name__ == 'RandomSampler': |
| 131 | data_sampler = data_sampler(dset, replacement, num_samples, generator) |
| 132 | else: |
| 133 | raise RuntimeError(f"{data_sampler.__name__} is not implemented.") |
| 134 | |
| 135 | if distributed: |
| 136 | ''' |
| 137 | Different with DistributedSampler, |
| 138 | the DistribuedProxySampler does not shuffle the data (just wrapper for dist). |
| 139 | ''' |
| 140 | data_sampler = DistributedProxySampler(data_sampler) |
| 141 | |
| 142 | batch_sampler = BatchSampler(data_sampler, batch_size, drop_last) |
no test coverage detected