Collate gives index out of range when iterating over dataloader in pytorch geometric

[LarsDoorenbos]

I am trying to use a dataloader with graphs I load in. When using this with toy data it works, but with the loaded data it does not. Both times the graphs are Data objects, the only difference is added attributes. Why does this raise an error in the collate function?

from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
import torch 
import pickle
import skimage
import torchvision
from skimage.future.graph import rag
from torch_geometric.utils import convert
import torchvision.transforms as transforms
import numpy as np

def save_obj(obj, name):
    with open('obj/'+ name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load_obj(name ):
    with open('obj/' + name + '.pkl', 'rb') as f:
        return pickle.load(f)

transform = transforms.ToTensor()

trainSet = torchvision.datasets.CIFAR10(root='../data', train=True, download=True, transform=transform)

img = trainSet[0][0].numpy().astype(float)
img = np.moveaxis(img, 0, 2)

seg = skimage.segmentation.slic(img, n_segments=70, start_label=0, convert2lab=True)
img_rag = rag.rag_mean_color(img, seg, mode='distance')
img_rag.merge_nodes(1, 2)
G1 = convert.from_networkx(img_rag, group_node_attrs=['mean color', 'pixel count'])

img = trainSet[1][0].numpy().astype(float)
img = np.moveaxis(img, 0, 2)

seg = skimage.segmentation.slic(img, n_segments=70, start_label=0, convert2lab=True)
img_rag = rag.rag_mean_color(img, seg, mode='distance')
img_rag.merge_nodes(1, 2)
G2 = convert.from_networkx(img_rag, group_node_attrs=['mean color', 'pixel count'])

data_list = [G1, G2]

save_obj(data_list, 'test')
data_list = load_obj('test')

print(G1)
print(G2)

loader = DataLoader(data_list, batch_size=32)
print(next(iter(loader)))

This gives:

Data(x=[3, 1], edge_index=[2, 4])
Data(x=[3, 1], edge_index=[2, 4])
Batch(x=[6, 1], edge_index=[2, 8], batch=[6], ptr=[3])
Data(edge_index=[2, 96], labels=[20], pixel count=[20], total color=[20, 3], mean color=[20, 3], weight=[96], num_nodes=20, x=[20, 4], y=0)
Data(edge_index=[2, 96], labels=[20], pixel count=[20], total color=[20, 3], mean color=[20, 3], weight=[96], num_nodes=20, x=[20, 4], y=0)
Traceback (most recent call last):
  File "test.py", line 39, in <module>
    print(next(iter(loader)))
  File "/home/-/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()
  File "/home/-/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 561, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/-/.local/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 47, in fetch
    return self.collate_fn(data)
  File "/home/-/anaconda3/lib/python3.8/site-packages/torch_geometric/loader/dataloader.py", line 39, in __call__
    return self.collate(batch)
  File "/home/-/anaconda3/lib/python3.8/site-packages/torch_geometric/loader/dataloader.py", line 19, in collate
    return Batch.from_data_list(batch, self.follow_batch,
  File "/home/-/anaconda3/lib/python3.8/site-packages/torch_geometric/data/batch.py", line 63, in from_data_list
    batch, slice_dict, inc_dict = collate(
  File "/home/-/anaconda3/lib/python3.8/site-packages/torch_geometric/data/collate.py", line 76, in collate
    value, slices, incs = _collate(attr, values, data_list, stores,
  File "/home/-/anaconda3/lib/python3.8/site-packages/torch_geometric/data/collate.py", line 127, in _collate
    value, slices, incs = _collate(key, [v[i] for v in values],
  File "/home/-/anaconda3/lib/python3.8/site-packages/torch_geometric/data/collate.py", line 127, in _collate
    value, slices, incs = _collate(key, [v[i] for v in values],
  File "/home/-/anaconda3/lib/python3.8/site-packages/torch_geometric/data/collate.py", line 127, in <listcomp>
    value, slices, incs = _collate(key, [v[i] for v in values],
IndexError: list index out of range

[rusty1s]

How do you save your graphs? The following works for me:

from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
import torch
import pickle


def load_obj(name):
    with open(name + '.pkl', 'rb') as f:
        return pickle.load(f)


# Toy data
edge_index = torch.tensor([[0, 1], [1, 0], [1, 2], [2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

graph1 = Data(x=x, edge_index=edge_index.t().contiguous())
graph2 = Data(x=x, edge_index=edge_index.t().contiguous())

print(graph1)
print(graph2)
data_list = [graph1, graph2]

loader = DataLoader(data_list, batch_size=32)
print(next(iter(loader)))

with open('train0.pkl', 'wb') as f:
    pickle.dump(data_list, f)

#Own data:
train_graphs1 = load_obj('train0')

graph1 = train_graphs1[0]
graph2 = train_graphs1[1]

print(graph1)
print(graph2)

data_list = [graph1, graph2]

loader = DataLoader(data_list, batch_size=32)
print(next(iter(loader)))

[rusty1s]

Mh, that works for me as well. Can you give some more information about your machine's setup? (Python, OS, CUDA, PyG). Furthermore, can you try to install PyG from master and see if that fixes your issues already?

pip install git+https://github.com/pyg-team/pytorch_geometric.git

[rusty1s]

Furthermore, were your Batch objects saved with PyG < 2.0?

[LarsDoorenbos]

Install PyG from master seems to have fixed the issues, thank you!

[star20174154]

Hello, I have the same problem where torch-geometric==2.0.1 . How did you solve it please?

[rusty1s]

You can try if upgrading PyG alone fixes the issue.