Class StreamSampler

imperative/python/megengine/data/sampler.py:148–175

r"""Sampler for stream dataset. Warning: In the case of multiple machines, sampler should ensure that each worker gets different data. But this class cannot do it yet, please build your own dataset and sampler to achieve this goal. Usually, :meth:`~.StreamDataset.__

class StreamSampler(Sampler):
    r"""Sampler for stream dataset.

    Warning:
        In the case of multiple machines, sampler should ensure that each worker gets
        different data. But this class cannot do it yet, please build your own
        dataset and sampler to achieve this goal.

    Usually, :meth:`~.StreamDataset.__iter__` can return different iterator by
    ``rank = dist.get_rank()``. So that they will get different data.
    """

    def __init__(self, batch_size=1):
        self.batch_size = batch_size

    def __iter__(self):
        return self.batch()

    def batch(self):
        batch = []
        for idx in self.sample():
            batch.append(idx)
            if len(batch) == self.batch_size:
                yield batch
                batch = []

    def sample(self):
        return count(start=0)


class SequentialSampler(MapSampler):