Correct way to create a NeighborLoader for heterogeneous graphs

[mtoshevska]

Hello!

I am trying to create a GNN-based model for node classification on a heterogeneous graph (the graph has 3 types of nodes and 4 types of edges while the goal is to classify nodes from one of the 3 types). Full-batch training of the models works fine. However, when I try mini-batch training with NeighborLoader for sampling, I get the following error: Can't hash IValues with tag 'String'.

The code for creating the NeighborLoader instance (sampling 8 neighbors for each edge type for each of the 2 GNN layers; dataset is an instance of HeteroData):
train_loader = NeighborLoader(dataset, num_neighbors={key: [8] * 2 for key in dataset.edge_types}, shuffle=True, input_nodes=train_input_nodes, batch_size=16)
Iterating through this loader or using it for training with PyTorch-Lightning throws the aforementioned error.

I would like to ask what am I doing wrong and what is the correct way for creating NeighborLoader instances for heterogeneous graphs?

[akihironitta]

I'm happy to help if you could provide the full error message and a script to reproduce it :)

[mtoshevska]

Since my original dataset is big, I have created a sample script with the OGB_MAG dataset that throws the same error:

import torch
import os.path as osp
import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.loader import NeighborLoader


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

if __name__ == '__main__':
    path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/OGB')
    transform = T.ToUndirected(merge=True)
    dataset = OGB_MAG(path, preprocess='metapath2vec', transform=transform)
    data = dataset[0].to(device, 'x', 'y')

    train_input_nodes = ('paper', data['paper'].train_mask)
    train_loader = NeighborLoader(data, num_neighbors={key: [2] * 2 for key in data.edge_types}, shuffle=True,
                                  input_nodes=train_input_nodes, batch_size=16)

    next(iter(train_loader))

This is the full error message:

Traceback (most recent call last):
  File "D:\Python projects\Text-Classification-GCN\gnn\mini-batch_heterodata.py", line 20, in <module>
    next(iter(train_loader))
  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch_geometric\loader\base.py", line 36, in __next__
    return self.transform_fn(next(self.iterator))
  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch\utils\data\dataloader.py", line 652, in __next__
    data = self._next_data()
  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch\utils\data\dataloader.py", line 692, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch\utils\data\_utils\fetch.py", line 52, in fetch
    return self.collate_fn(data)
  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch_geometric\loader\node_loader.py", line 112, in collate_fn
    out = self.node_sampler.sample_from_nodes(input_data)
  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch_geometric\sampler\neighbor_sampler.py", line 336, in sample_from_nodes
    return node_sample(index, self._sample, self.input_type, **kwargs)
  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch_geometric\sampler\neighbor_sampler.py", line 376, in node_sample
    seed_time_dict=seed_time_dict)
  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch_geometric\sampler\neighbor_sampler.py", line 267, in _sample
    self.directed,
  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch\_ops.py", line 143, in __call__
    return self._op(*args, **kwargs or {})
RuntimeError: Can't hash IValues with tag 'String'

[akihironitta]

Which versions are you using? I might not be of any help here, but the following line looks related to torch_sparse:

  File "C:\Users\mtosh\AppData\Local\Programs\Python\Python37\lib\site-packages\torch_geometric\sampler\neighbor_sampler.py", line 267, in _sample
    self.directed,

Just for the note, I have no issue running the script in my Linux env:

torch                    2.0.1+cu118
torch-cluster            1.6.1+pt20cu118
torch_geometric          2.4.0            /workspaces/pytorch_geometric
torch-scatter            2.1.1+pt20cu118
torch-sparse             0.6.17+pt20cu118

[mtoshevska]

It appears that the problem was in the versions and probably the OS as well (I was testing the code on Windows 11 with Python 3.7 and the following versions:
torch - 1.12.0
torch-cluster - 1.6.0
torch-geometric - 2.2.0
torch-scatter - 2.0.9
torch-sparse - 0.6.13

When I run the code on a Unix-based OS with newer versions of the libraries it does not throw any error.