EdgeConv - IndexError: index out of range in self

[manami-16]

Hello,
I have been trying to create a EdgeConv model by converting an image dataset of 32*32 to a graph using KNNGraph.
However, I got IndexError: index out of range in self, which I have been stuck on.
I would really appreciate it if anybody can figure out what the issue is and how I can resolve it.

Set up the dataset

## random_tensor will be replaced to actual image dataset
random_tensor = torch.rand((400, 1, 32, 32))
random_tensor = DataLoader(random_tensor, batch_size=32, shuffle=False, drop_last=True)

Define a function that returns edge_index and position from KNNGraph

def get_grid(height=32, width=32):
    edge_index = []
    pos = []
    for data in batch:
        edge_idx, position = grid(height, width)
        edge_index.append(edge_idx)
        pos.append(position)
        
    edge_index = torch.cat(edge_index)
    pos = torch.cat(pos)

    return edge_index, pos

Define nn in EdgeConv

## in_channel = 1 because input=(batch_size, 1, 32, 32)

layer = nn.Sequential(nn.Conv2d(1, 5, kernel_size=2, stride=1, padding=1), 
                      nn.ReLU()).double()
edge_conv = EdgeConv(nn=layer, aggr='max', node_dim=0)

Pass each batch to get knn_data

for batch in random_tensor:
    edge_index, pos = get_grid(32, 32)
    data = Data(x=batch, edge_index=edge_index, pos=pos)
    
    knn = KNNGraph(k=3)
    knn_data = knn(data)

print(knn_data) -> Data(x=[1, 32, 32], edge_index=[2, 3072], pos=[1024, 2])

Feed the data into edge_conv

x = edge_conv(knn_data.x, knn_data.edge_index)

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Input In [15], in <cell line: 1>()
----> 1 x = edge_conv(knn_data.x, knn_data.edge_index)

File ~/miniconda3/envs/gnnpyg/lib/python3.8/site-packages/torch/nn/modules/module.py:1110, in Module._call_impl(self, *input, **kwargs)
   1106 # If we don't have any hooks, we want to skip the rest of the logic in
   1107 # this function, and just call forward.
   1108 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110     return forward_call(*input, **kwargs)
   1111 # Do not call functions when jit is used
   1112 full_backward_hooks, non_full_backward_hooks = [], []

File ~/miniconda3/envs/gnnpyg/lib/python3.8/site-packages/torch_geometric/nn/conv/edge_conv.py:61, in EdgeConv.forward(self, x, edge_index)
     59     x: PairTensor = (x, x)
     60 # propagate_type: (x: PairTensor)
---> 61 return self.propagate(edge_index, x=x, size=None)

File ~/miniconda3/envs/gnnpyg/lib/python3.8/site-packages/torch_geometric/nn/conv/message_passing.py:309, in MessagePassing.propagate(self, edge_index, size, **kwargs)
    306     for arg in decomp_args:
    307         kwargs[arg] = decomp_kwargs[arg][i]
--> 309 coll_dict = self.__collect__(self.__user_args__, edge_index,
    310                              size, kwargs)
    312 msg_kwargs = self.inspector.distribute('message', coll_dict)
    313 for hook in self._message_forward_pre_hooks.values():

File ~/miniconda3/envs/gnnpyg/lib/python3.8/site-packages/torch_geometric/nn/conv/message_passing.py:202, in MessagePassing.__collect__(self, args, edge_index, size, kwargs)
    200         if isinstance(data, Tensor):
    201             self.__set_size__(size, dim, data)
--> 202             data = self.__lift__(data, edge_index, dim)
    204         out[arg] = data
    206 if isinstance(edge_index, Tensor):

File ~/miniconda3/envs/gnnpyg/lib/python3.8/site-packages/torch_geometric/nn/conv/message_passing.py:172, in MessagePassing.__lift__(self, src, edge_index, dim)
    170 if isinstance(edge_index, Tensor):
    171     index = edge_index[dim]
--> 172     return src.index_select(self.node_dim, index)
    173 elif isinstance(edge_index, SparseTensor):
    174     if dim == 1:

IndexError: index out of range in self

Just for reference,

knn_data.edge_index.max() < knn_data.x.size(0)
>> False

[rusty1s]

You need to reshape your image data to [32*32, 1] - it now becomes a node feature matrix :)

[rusty1s]

At best, you reshape per example before batching takes place. After batching, you will receive a [batch_size*32*32, 1] matrix.

[rusty1s]

Here is an example: https://gist.github.com/rusty1s/866ca8b830f22aa34923a0f8164f6e64

[manami-16]

Thank you so much! I could resolve the error!

[LukeLIN-web]

I meet the same problem in SageConv

[rusty1s]

Can you clarify? If the index_select call throws an error, this usually indicates that indices in edge_index point to nodes outside the [0, num_nodes) range.