Class ArrayDataset

imperative/python/megengine/data/dataset/meta_dataset.py:109–156

r"""ArrayDataset is a dataset for numpy array data. One or more numpy arrays are needed to initiate the dataset. And the dimensions represented sample number are expected to be the same. Args: Arrays(dataset and labels): the datas and labels to be returned iteratively. Ret

class ArrayDataset(Dataset):
    r"""ArrayDataset is a dataset for numpy array data.

    One or more numpy arrays are needed to initiate the dataset.
    And the dimensions represented sample number are expected to be the same.

    Args:
        Arrays(dataset and labels): the datas and labels to be returned iteratively.

    Returns:
        Tuple: A set of raw data and corresponding label.


    Examples:

        .. code-block:: python

            from megengine.data.dataset import ArrayDataset
            from megengine.data.dataloader import DataLoader
            from megengine.data.sampler import SequentialSampler

            rand_data = np.random.randint(0, 255, size=(sample_num, 1, 32, 32), dtype=np.uint8)
            label = np.random.randint(0, 10, size=(sample_num,), dtype=int)
            dataset = ArrayDataset(rand_data, label)
            seque_sampler = SequentialSampler(dataset, batch_size=2)

            dataloader = DataLoader(
                dataset,
                sampler = seque_sampler,
                num_workers=3,
            )

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

    """

    def __init__(self, *arrays):
        super().__init__()
        if not all(len(arrays[0]) == len(array) for array in arrays):
            raise ValueError("lengths of input arrays are inconsistent")
        self.arrays = arrays

    def __getitem__(self, index: int) -> Tuple:
        return tuple(array[index] for array in self.arrays)

    def __len__(self) -> int:
        return len(self.arrays[0])


class ConcatDataset(Dataset):