GraphGym: graph classification with PyG datasets raises KeyError: 'train_graph_index'

🐛 Describe the bug

When setting up a configuration for graph classification with a dataset in PyG format an error message is received during loading of the dataset.

Error message:

Traceback (most recent call last):
  File "main.py", line 38, in <module>
    loaders = create_loader()
  File ".../lib/python3.8/site-packages/torch_geometric/graphgym/loader.py", line 301, in create_loader
    id = dataset.data['train_graph_index']
  File ".../lib/python3.8/site-packages/torch_geometric/data/data.py", line 371, in __getitem__
    return self._store[key]
  File ".../lib/python3.8/site-packages/torch_geometric/data/storage.py", line 67, in __getitem__
    return self._mapping[key]
KeyError: 'train_graph_index'

This is the configuration file I use:

out_dir: results
dataset:
  format: PyG
  name: MNIST
  task: graph
  task_type: classification
  transductive: false

The error appears to be that the attribute “train_graph_index” is only set for OGB-formatted datasets in load_ogb() in graphgym/loader.py, but is expected to be present for every dataset in create_loader.

Environment

• PyG version: 2.0.4
• PyTorch version: 1.10.2
• OS: Linux, kernel 5.10
• Python version: 3.8.12
• CUDA/cuDNN version: ./.
• How you installed PyTorch and PyG (conda, pip, source): conda
• Any other relevant information (e.g., version of torch-scatter):

[HelgeS]

Yes, using a random split would work, but I understand we should use a given split if available.

I only picked MNIST as an example here, but ultimately want to use my own dataset in PyG format; there a random split would be fine.

[rusty1s]

@JiaxuanYou @HelgeS Any updates on this?