Sampler that restricts data loading to a subset of the dataset. Modified from torch.utils.data.distributed.DistributedSampler Support enlarging the dataset for iteration-based training, for saving time when restart the dataloader after each epoch Args: dataset (torch.utils.
| 4 | |
| 5 | |
| 6 | class EnlargedSampler(Sampler): |
| 7 | """Sampler that restricts data loading to a subset of the dataset. |
| 8 | |
| 9 | Modified from torch.utils.data.distributed.DistributedSampler |
| 10 | Support enlarging the dataset for iteration-based training, for saving |
| 11 | time when restart the dataloader after each epoch |
| 12 | |
| 13 | Args: |
| 14 | dataset (torch.utils.data.Dataset): Dataset used for sampling. |
| 15 | num_replicas (int | None): Number of processes participating in |
| 16 | the training. It is usually the world_size. |
| 17 | rank (int | None): Rank of the current process within num_replicas. |
| 18 | ratio (int): Enlarging ratio. Default: 1. |
| 19 | """ |
| 20 | |
| 21 | def __init__(self, dataset, num_replicas, rank, ratio=1): |
| 22 | self.dataset = dataset |
| 23 | self.num_replicas = num_replicas |
| 24 | self.rank = rank |
| 25 | self.epoch = 0 |
| 26 | print(self.dataset,ratio,self.num_replicas) |
| 27 | self.num_samples = math.ceil(len(self.dataset) * ratio / self.num_replicas) |
| 28 | self.total_size = self.num_samples * self.num_replicas |
| 29 | |
| 30 | def __iter__(self): |
| 31 | # deterministically shuffle based on epoch |
| 32 | g = torch.Generator() |
| 33 | g.manual_seed(self.epoch) |
| 34 | indices = torch.randperm(self.total_size, generator=g).tolist() |
| 35 | |
| 36 | dataset_size = len(self.dataset) |
| 37 | indices = [v % dataset_size for v in indices] |
| 38 | |
| 39 | # subsample |
| 40 | indices = indices[self.rank:self.total_size:self.num_replicas] |
| 41 | assert len(indices) == self.num_samples |
| 42 | |
| 43 | return iter(indices) |
| 44 | |
| 45 | def __len__(self): |
| 46 | return self.num_samples |
| 47 | |
| 48 | def set_epoch(self, epoch): |
| 49 | self.epoch = epoch |
no outgoing calls
no test coverage detected