Stratified K Fold and Dataset

🐛 Bug

Hi there! The Pytorch geometric Dataset object used to work nicely with scikit-learn’s StratifiedKFold. See an example below:

 kf2 = StratifiedKFold(n_splits=9, shuffle=False)
 for train_idxs, val_idxs in kf2.split(dataset, dataset.data.y):
     train_dataset = dataset[torch.LongTensor(train_idx)]

In version 1.3.2 the dataset slicing used to also affect dataset.data.y. However, slicing now seems to leave this unchanged. So if one wants to further split train_dataset using StratifiedKFold, this is not working anymore because len(train_dataset) and len(train_dataset.data.y) won’t match.

To Reproduce

Steps to reproduce the behavior:

1. Load Proteins dataset from PyTorch geometric
2. Check data len(dataset[indices]) is different from len(dataset[indices].data.y).

Is this the new expected behaviour and is there any elegant workaround for this?

I think the older behaviour from version 1.3.2 was more natural.

Environment

• OS: MacOS
• Python version: 3.6
• PyTorch version: 1.5
• CUDA/cuDNN version: cpu

[rusty1s]

Yes, that is correct. This should be already fixed in later versions.

[rusty1s]

I just added a dataset.copy() method that converts sliced datasets back to a contiguous layout:

from torch_geometric.datasets import TUDataset

dataset = TUDataset('/tmp/TUDataset', 'MUTAG')
print(len(dataset))
>>> 180
print(dataset.data.y.size())
>>> torch.Size([188])

dataset = dataset[:40]
print(len(dataset))
>>> 40
print(dataset.data.y.size())  # <- invalid access
>>> torch.Size([188])

dataset = dataset.copy()
print(len(dataset))
>>> 40
print(dataset.data.y.size())  # <- correct access
>>> torch.Size([40])