Create Dataset class that accept variable number of data arrays?

1

Hi, currently I have implemented following Dataset class that accepts input and output data for creating batches for training.

    class My_dataset(Dataset):
        def __init__(self, data_in, data_out):
            self._data_in = data_in
            self._data_out = data_out

        def __len__(self):
            return len(self._data_in)

        def __getitem__(self, index):
            return self._data_in[index], self._data_out[index]

Now I would like to add a few helper variables that has the same length with data to improve training results and they should also be split into same batches with input and output data. Does anyone know if it is possible to create a subclass of Dataset that accepts variable number of arrays in the constructor (e.g. it could be called as My_dataset(data_in, data_out, helper_1), My_dataset(data_in, data_out, helper_1, helper_2), etc) and generates batches accordingly?

Thank you!

2

Never mind, it is solved:

    class My_dataset(Dataset):
        def __init__(self, *data):
            # accept variable number of arrays to construct dataset
            self._data = data

        def __len__(self):
            return len(self._data[0])

        def __getitem__(self, index):
            return [item[index] for item in self._data]
3

Hi,

Note that this is actually already implemented in TensorDataset, see the implementation here