import numpy as np
from paddle.io import Dataset

# define a random dataset
class RandomDataset(Dataset):
    def __init__(self, num_samples):
        self.num_samples = num_samples

    def __iter__(self):
        for i in range(self.num_samples):
            image = np.random.random([784]).astype('float32')
            label = np.random.randint(0, 9, (1, )).astype('int64')
            yield image, label

dataset = RandomDataset(10)
for img, lbl in dataset:
    print(img, lbl)

import math
import paddle
import numpy as np
from paddle.io import IterableDataset, DataLoader, get_worker_info

class SplitedIterableDataset(IterableDataset):
    def __init__(self, start, end):
        self.start = start
        self.end = end

    def __iter__(self):
        worker_info = get_worker_info()
        if worker_info is None:
            iter_start = self.start
            iter_end = self.end
        else:
            per_worker = int(
                math.ceil((self.end - self.start) / float(
                    worker_info.num_workers)))
            worker_id = worker_info.id
            iter_start = self.start + worker_id * per_worker
            iter_end = min(iter_start + per_worker, self.end)

        for i in range(iter_start, iter_end):
            yield np.array([i])

dataset = SplitedIterableDataset(start=2, end=9)
dataloader = DataLoader(
    dataset,
    num_workers=2,
    batch_size=1,
    drop_last=True)

for data in dataloader:
    print(data)
# outputs: [2, 5, 3, 6, 4, 7]

import math
import paddle
import numpy as np
from paddle.io import IterableDataset, DataLoader, get_worker_info

class RangeIterableDataset(IterableDataset):
    def __init__(self, start, end):
        self.start = start
        self.end = end

    def __iter__(self):
        for i in range(self.start, self.end):
            yield np.array([i])

dataset = RangeIterableDataset(start=2, end=9)

def worker_init_fn(worker_id):
    worker_info = get_worker_info()

    dataset = worker_info.dataset
    start = dataset.start
    end = dataset.end
    num_per_worker = int(
        math.ceil((end - start) / float(worker_info.num_workers)))

    worker_id = worker_info.id
    dataset.start = start + worker_id * num_per_worker
    dataset.end = min(dataset.start + num_per_worker, end)

dataloader = DataLoader(
    dataset,
    num_workers=2,
    batch_size=1,
    drop_last=True,
    worker_init_fn=worker_init_fn)

for data in dataloader:
    print(data)
# outputs: [2, 5, 3, 6, 4, 7]