Class SequentialSampler

imperative/python/megengine/data/sampler.py:178–215

r"""Sample elements sequentially. Args: dataset: dataset to sample from. batch_size: batch size for batch method. drop_last: set ``True`` to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If ``False`` and t

class SequentialSampler(MapSampler):
    r"""Sample elements sequentially.

    Args:
        dataset: dataset to sample from.
        batch_size: batch size for batch method.
        drop_last: set ``True`` to drop the last incomplete batch,
            if the dataset size is not divisible by the batch size. If ``False`` and
            the size of dataset is not divisible by the batch_size, then the last batch will
            be smaller. Default: False
        indices: indice of samples.
        world_size: number of ranks.
        rank: rank id, non-negative interger within 0 and ``world_size``.
    """

    def __init__(
        self,
        dataset,
        batch_size=1,
        drop_last=False,
        indices=None,
        world_size=None,
        rank=None,
    ):
        super().__init__(dataset, batch_size, drop_last, None, world_size, rank)
        if indices is not None and not isinstance(indices, collections.abc.Sequence):
            raise ValueError(
                "indices should be None or a sequence, "
                "but got indices={}".format(indices)
            )
        self.indices = indices

    def sample(self) -> Iterator[Any]:
        r"""Return a generator."""
        if self.indices is None:
            return iter(range(len(self.dataset)))
        else:
            return self.indices


class RandomSampler(MapSampler):