faisalhr1997/codeformer

1import math2import torch3from torch.utils.data.sampler import Sampler4 5 6class EnlargedSampler(Sampler):7    """Sampler that restricts data loading to a subset of the dataset.8 9    Modified from torch.utils.data.distributed.DistributedSampler10    Support enlarging the dataset for iteration-based training, for saving11    time when restart the dataloader after each epoch12 13    Args:14        dataset (torch.utils.data.Dataset): Dataset used for sampling.15        num_replicas (int | None): Number of processes participating in16            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 = dataset23        self.num_replicas = num_replicas24        self.rank = rank25        self.epoch = 026        self.num_samples = math.ceil(len(self.dataset) * ratio / self.num_replicas)27        self.total_size = self.num_samples * self.num_replicas28 29    def __iter__(self):30        # deterministically shuffle based on epoch31        g = torch.Generator()32        g.manual_seed(self.epoch)33        indices = torch.randperm(self.total_size, generator=g).tolist()34 35        dataset_size = len(self.dataset)36        indices = [v % dataset_size for v in indices]37 38        # subsample39        indices = indices[self.rank:self.total_size:self.num_replicas]40        assert len(indices) == self.num_samples41 42        return iter(indices)43 44    def __len__(self):45        return self.num_samples46 47    def set_epoch(self, epoch):48        self.epoch = epoch49