Distributed Sampler that subsamples indicies sequentially, making it easier to collate all results at the end. Even though we only use this sampler for eval and predict (no training), which means that the model params won't have to be synced (i.e. will not hang for synchronizat
| 84 | |
| 85 | |
| 86 | class SequentialDistributedSampler(Sampler): |
| 87 | """ |
| 88 | Distributed Sampler that subsamples indicies sequentially, |
| 89 | making it easier to collate all results at the end. |
| 90 | |
| 91 | Even though we only use this sampler for eval and predict (no training), |
| 92 | which means that the model params won't have to be synced (i.e. will not hang |
| 93 | for synchronization even if varied number of forward passes), we still add extra |
| 94 | samples to the sampler to make it evenly divisible (like in `DistributedSampler`) |
| 95 | to make it easy to `gather` or `reduce` resulting tensors at the end of the loop. |
| 96 | """ |
| 97 | |
| 98 | def __init__(self, dataset, num_replicas=None, rank=None): |
| 99 | if num_replicas is None: |
| 100 | if not torch.distributed.is_available(): |
| 101 | raise RuntimeError("Requires distributed package to be available") |
| 102 | num_replicas = torch.distributed.get_world_size() |
| 103 | if rank is None: |
| 104 | if not torch.distributed.is_available(): |
| 105 | raise RuntimeError("Requires distributed package to be available") |
| 106 | rank = torch.distributed.get_rank() |
| 107 | self.dataset = dataset |
| 108 | self.num_replicas = num_replicas |
| 109 | self.rank = rank |
| 110 | self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas)) |
| 111 | self.total_size = self.num_samples * self.num_replicas |
| 112 | |
| 113 | def __iter__(self): |
| 114 | indices = list(range(len(self.dataset))) |
| 115 | |
| 116 | # add extra samples to make it evenly divisible |
| 117 | indices += indices[: (self.total_size - len(indices))] |
| 118 | assert len(indices) == self.total_size |
| 119 | |
| 120 | # subsample |
| 121 | indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples] |
| 122 | assert len(indices) == self.num_samples |
| 123 | |
| 124 | return iter(indices) |
| 125 | |
| 126 | def __len__(self): |
| 127 | return self.num_samples |
| 128 | |
| 129 | |
| 130 | def get_tpu_sampler(dataset: Dataset): |
no outgoing calls
no test coverage detected