Sampler that restricts data loading to a subset of the dataset. It is especially useful in conjunction with :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each process can pass a DistributedSampler instance as a DataLoader sampler, and load a subset of the origina
| 7 | from PIL import Image, ImageEnhance, ImageOps |
| 8 | |
| 9 | class SubsetDistributedSampler(Sampler): |
| 10 | """Sampler that restricts data loading to a subset of the dataset. |
| 11 | |
| 12 | It is especially useful in conjunction with |
| 13 | :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each |
| 14 | process can pass a DistributedSampler instance as a DataLoader sampler, |
| 15 | and load a subset of the original dataset that is exclusive to it. |
| 16 | |
| 17 | .. note:: |
| 18 | Dataset is assumed to be of constant size. |
| 19 | |
| 20 | Arguments: |
| 21 | dataset: Dataset used for sampling. |
| 22 | num_replicas (optional): Number of processes participating in |
| 23 | distributed training. |
| 24 | rank (optional): Rank of the current process within num_replicas. |
| 25 | shuffle (optional): If true (default), sampler will shuffle the indices |
| 26 | """ |
| 27 | |
| 28 | def __init__(self, dataset, indices, num_replicas=None, rank=None, shuffle=True): |
| 29 | if num_replicas is None: |
| 30 | if not dist.is_available(): |
| 31 | raise RuntimeError("Requires distributed package to be available") |
| 32 | num_replicas = dist.get_world_size() |
| 33 | if rank is None: |
| 34 | if not dist.is_available(): |
| 35 | raise RuntimeError("Requires distributed package to be available") |
| 36 | rank = dist.get_rank() |
| 37 | self.dataset = dataset |
| 38 | self.num_replicas = num_replicas |
| 39 | self.rank = rank |
| 40 | self.epoch = 0 |
| 41 | self.indices = indices |
| 42 | self.num_samples = int(math.ceil(len(self.indices) * 1.0 / self.num_replicas)) |
| 43 | self.total_size = self.num_samples * self.num_replicas |
| 44 | self.shuffle = shuffle |
| 45 | |
| 46 | def __iter__(self): |
| 47 | # deterministically shuffle based on epoch |
| 48 | g = torch.Generator() |
| 49 | g.manual_seed(self.epoch) |
| 50 | if self.shuffle: |
| 51 | # indices = torch.randperm(len(self.dataset), generator=g).tolist() |
| 52 | indices = list(self.indices[i] for i in torch.randperm(len(self.indices))) |
| 53 | else: |
| 54 | # indices = list(range(len(self.dataset))) |
| 55 | indices = self.indices |
| 56 | |
| 57 | |
| 58 | # add extra samples to make it evenly divisible |
| 59 | indices += indices[:(self.total_size - len(indices))] |
| 60 | assert len(indices) == self.total_size |
| 61 | |
| 62 | # subsample |
| 63 | indices = indices[self.rank:self.total_size:self.num_replicas] |
| 64 | assert len(indices) == self.num_samples |
| 65 | |
| 66 | return iter(indices) |
no outgoing calls
no test coverage detected