I have a Pandas dataframe with n
rows and k
columns loaded into memory. I would like to get batches for a forecasting task where the first training example of a batch should have shape (q, k)
with q
referring to the number of rows from the original dataframe (e.g. 0:128). The next example should be (128:256, k)
and so on. So, ultimately, one batch should have the shape (32, q, k)
with 32 corresponding to the batch size.
Since TensorDataset
from data_utils
does not work here, I am wondering what the best way would be. I tried to use np.array_split()
to get as first dimension the number of possible splits of q values in order to write a custom DataLoader but then reshaping is not guaranteed to work since not all arrays have the same shape.
Here is a minimal example to make it more clear. In this case, batch size is 3 and q is 2:
import pandas as pd import numpy as np df = pd.DataFrame(data=np.arange(0,30).reshape(10,3),columns=['A','B','C'])
The dataset:
A B C 0 0 1 2 1 3 4 5 2 6 7 8 3 9 10 11 4 12 13 14 5 15 16 17 6 18 19 20 7 21 22 23 8 24 25 26 9 27 28 29
The first batch in this case should have the shape (3,2,3) and look like:
array([[[ 0., 1., 2.], [ 3., 4., 5.]], [[ 3., 4., 5.], [ 6., 7., 8.]], [[ 6., 7., 8.], [ 9., 10., 11.]]])
Advertisement
Answer
You can write your analog of the TensorDataset. To do this you need to inherit from the Dataset class.
from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, data_frame, q): self.data = data_frame.values self.q = q def __len__(self): return self.data.shape[0] // self.q def __getitem__(self, index): return self.data[index * self.q: (index+1) * self.q]