Skip to content

infinite_random_sampler

data.sampler.infinite_random_sampler

InfiniteRandomSampler

Bases: RandomSampler

Infinitely samples elements randomly. If without replacement, then sample from a shuffled dataset. If with replacement, then user can specify :attr:num_samples to draw.

Parameters:

  • data_source (Dataset) –

    dataset to sample from

  • replacement (bool) –

    samples are drawn on-demand with replacement if True, default=False

  • num_samples (int) –

    number of samples to draw, default=len(dataset).

  • generator (Generator) –

    Generator used in sampling.

Assume that this DataLoader is used for training with shuffle=True.

We use InfiniteRandomSampler to make the DataLoader to infinitely sample the dataset. This prevents the "drop_last" and the "prefeching" problem across epochs. Be careful that you should stop the training loop by counting the number of steps manually. The DataLoader will never stop by itself.

NOTE: Note that this sampler has infinite length and thus you should be careful when calculating the current epoch. We recommend you to use step//steps_per_epoch to calculate the current epoch.

__iter__ 

__iter__() -> Iterator[int]
Source code in SaigeToolkit/data/sampler/infinite_random_sampler.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __iter__(self) -> Iterator[int]:
    n = len(self.data_source)
    if self.generator is None:
        seed = int(torch.empty((), dtype=torch.int64).random_().item())
        rng = torch.Generator()
        rng.manual_seed(seed)
    else:
        rng = self.generator

    while True:
        if self.replacement:
            yield from torch.randint(high=n, size=(32,), dtype=torch.int64, generator=rng).tolist()
        else:
            yield from torch.randperm(n, generator=rng).tolist()

__len__ 

__len__() -> int
Source code in SaigeToolkit/data/sampler/infinite_random_sampler.py 
41
42
def __len__(self) -> int:
    return 100_000_000_000_000