StreamDataset#

class StreamDataset[源代码]#

An abstract class for stream data.

__iter__ method is aditionally needed.

实际案例

from megengine.data.dataset import StreamDataset
from megengine.data.dataloader import DataLoader, get_worker_info
from megengine.data.sampler import StreamSampler

class MyStream(StreamDataset):
    def __init__(self):
        self.data = [iter([1, 2, 3]), iter([4, 5, 6]), iter([7, 8, 9])]
    def __iter__(self):
        worker_info = get_worker_info()
        data_iter = self.data[worker_info.idx]
        while True:
            yield next(data_iter)

dataloader = DataLoader(
    dataset = MyStream(),
    sampler = StreamSampler(batch_size=2),
    num_workers=3,
    parallel_stream = True,
)

for step, data in enumerate(dataloader):
    print(data)