I am trying to do a regression task on some graphs.
The data I have is structured like this:
- Tensor of floats that stores all of the node features “x” of shape (number of nodes, number of node features)
- Tensor of all edges “edge_index” that stores the indices of start and end nodes for each edge of shape (2, number of edges)
I want to be able to combine them into a tensor of shape (number of edges, 2*number of node features), i.e. combine the features of the nodes for every edge into a single tensor.
For example:
x = tensor([[ 0, 1, 2], #4 nodes with 3 features each
[ 3, 4, 5],
[ 6, 7, 8],
[ 9, 10, 11]])
edge_index = tensor([[0, 1, 2, 3], #4 edges. First edge goes from node0 to node1
[1, 2, 3, 0]]) Last edge goes from node3 to node0
result = tensor([[ 0, 1, 2, 3, 4, 5], #The combined tensor of shape (number of edges, 2*number of node features)
[ 3, 4, 5, 6, 7, 8],
[ 6, 7, 8, 9, 10, 11],
[ 9, 10, 11, 0, 1, 2]])
My first approach was to use torch.index_select, however after implementing it I realized that it doesn’t support backpropagation, therefore to me it was absolutely useless.
My second approach relied on slices:
output = output = torch.zeros([len(edge_index[0]),6], requires_grad = True)
for i in range(len(edge_index[0])):
output[i]=torch.cat((x[edge_index[0][i].item():edge_index[0][i].item()+1],x[edge_index[1][i].item():edge_index[1][i].item()+1]),1)
This approach gave me the desired result and even supports backpropagation, however it is extremely slow. Processing 1500 nodes with 8000 edges takes nearly 4.5 seconds!
I want to find an approach that both supports backpropagation and is fast.
Any suggestions would be greatly appreciated!
Thank you in advance!