| import math
|
| import torch
|
| from torch.utils.data.sampler import Sampler
|
|
|
|
|
| class EnlargedSampler(Sampler):
|
| """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.data.Dataset): Dataset used for sampling.
|
| num_replicas (int | None): Number of processes participating in
|
| the training. It is usually the world_size.
|
| rank (int | None): Rank of the current process within num_replicas.
|
| ratio (int): Enlarging ratio. Default: 1.
|
| """
|
|
|
| def __init__(self, dataset, num_replicas, rank, ratio=1):
|
| self.dataset = dataset
|
| self.num_replicas = num_replicas
|
| self.rank = rank
|
| self.epoch = 0
|
| self.num_samples = math.ceil(len(self.dataset) * ratio / self.num_replicas)
|
| self.total_size = self.num_samples * self.num_replicas
|
|
|
| def __iter__(self):
|
|
|
| g = torch.Generator()
|
| g.manual_seed(self.epoch)
|
| indices = torch.randperm(self.total_size, generator=g).tolist()
|
|
|
| dataset_size = len(self.dataset)
|
| indices = [v % dataset_size for v in indices]
|
|
|
|
|
| indices = indices[self.rank:self.total_size:self.num_replicas]
|
| assert len(indices) == self.num_samples
|
|
|
| return iter(indices)
|
|
|
| def __len__(self):
|
| return self.num_samples
|
|
|
| def set_epoch(self, epoch):
|
| self.epoch = epoch
|
|
|
