Skip to content

builder

data.sampler.builder

logger module-attribute 

logger = getLogger('SaigeResearch')

SAMPLER_IMPLEMENTATION module-attribute 

SAMPLER_IMPLEMENTATION = {'InfiniteRandomSampler': InfiniteRandomSampler}

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

build_sampler 

build_sampler(_target_: Optional[str], data_source: Optional[Sized], **cfg_sampler: dict) -> Optional[Sampler]

Parameters:

  • cfg_sampler (dict, default: {} ) –

    config dict for building sampler

Returns:

  • Optional[Sampler] –

    Optional[Sampler]: sampler object instance. if cfg_sampler is None, return None

Notes

Follows the interface of torch.utils.data.Sampler.

Source code in SaigeToolkit/data/sampler/builder.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def build_sampler(
    _target_: Optional[str], data_source: Optional[Sized], **cfg_sampler: dict
) -> Optional[Sampler]:
    """
    Args:
        cfg_sampler (dict): config dict for building sampler

    Returns:
        Optional[Sampler]: sampler object instance.
            if cfg_sampler is None, return None

    Notes:
        Follows the interface of torch.utils.data.Sampler.
    """
    if _target_ is None:
        return None
    sampler_class = get_sampler_class(_target_)
    logger.info(f"[{'SAMPLER'.center(9)}] {_target_} [params] {cfg_sampler}")
    return sampler_class(data_source=data_source, **cfg_sampler)

get_sampler_class 

get_sampler_class(cfg_sampler_name: str) -> Type[Sampler]
Source code in SaigeToolkit/data/sampler/builder.py 
36
37
38
39
40
41
42
43
44
45
def get_sampler_class(cfg_sampler_name: str) -> Type[Sampler]:
    try:
        return SAMPLER_IMPLEMENTATION[cfg_sampler_name]
    except KeyError as e:
        raise RuntimeError(
            f"Sampler {cfg_sampler_name} not implemented. Supported sampler list: {SAMPLER_IMPLEMENTATION.keys()}"
        ) from e
    except Exception as e:
        logger.error(f"Error: {e}")
        raise e