|
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 |
|
|
