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 original
| 36 | |
| 37 | |
| 38 | class NodeDistributedSampler(Sampler): |
| 39 | """Sampler that restricts data loading to a subset of the dataset. |
| 40 | It is especially useful in conjunction with |
| 41 | :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each |
| 42 | process can pass a DistributedSampler instance as a DataLoader sampler, |
| 43 | and load a subset of the original dataset that is exclusive to it. |
| 44 | .. note:: |
| 45 | Dataset is assumed to be of constant size. |
| 46 | Arguments: |
| 47 | dataset: Dataset used for sampling. |
| 48 | num_replicas (optional): Number of processes participating in |
| 49 | distributed training. |
| 50 | rank (optional): Rank of the current process within num_replicas. |
| 51 | """ |
| 52 | |
| 53 | def __init__(self, |
| 54 | dataset, |
| 55 | num_replicas=None, |
| 56 | rank=None, |
| 57 | local_rank=None, |
| 58 | local_size=None): |
| 59 | if num_replicas is None: |
| 60 | if not dist.is_available(): |
| 61 | raise RuntimeError( |
| 62 | 'Requires distributed package to be available') |
| 63 | num_replicas = dist.get_world_size() |
| 64 | if rank is None: |
| 65 | if not dist.is_available(): |
| 66 | raise RuntimeError( |
| 67 | 'Requires distributed package to be available') |
| 68 | rank = dist.get_rank() |
| 69 | if local_rank is None: |
| 70 | local_rank = int(os.environ.get('LOCAL_RANK', 0)) |
| 71 | if local_size is None: |
| 72 | local_size = int(os.environ.get('LOCAL_SIZE', 1)) |
| 73 | self.dataset = dataset |
| 74 | self.num_replicas = num_replicas |
| 75 | self.num_parts = local_size |
| 76 | self.rank = rank |
| 77 | self.local_rank = local_rank |
| 78 | self.epoch = 0 |
| 79 | self.num_samples = int( |
| 80 | math.ceil(len(self.dataset) * 1.0 / self.num_replicas)) |
| 81 | self.total_size = self.num_samples * self.num_replicas |
| 82 | |
| 83 | self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts |
| 84 | |
| 85 | def __iter__(self): |
| 86 | # deterministically shuffle based on epoch |
| 87 | g = torch.Generator() |
| 88 | g.manual_seed(self.epoch) |
| 89 | |
| 90 | t = torch.Generator() |
| 91 | t.manual_seed(0) |
| 92 | |
| 93 | indices = torch.randperm(len(self.dataset), generator=t).tolist() |
| 94 | # indices = range(len(self.dataset)) |
| 95 | indices = [i for i in indices if i % self.num_parts == self.local_rank] |