How to combine the output results of two branch graph networks where each branch has a different number of nodes

[dhorka]

Hi,

I'm working in a network that is composed by two graph conv branches. Branch one has as input (N, F) where N is the number of nodes and F number of features and outputs a graph (N, F'). The second branch is similar to the first one, but it has a different number of input nodes (M, F) and outputs (M, F'). After both branches I have a third branch that is also a graph conv branch that combines the outputs of both previous branches as it can be seen in the following picture:

Currently what I am doing to combine both branches, is concatenate the node features, pos and batch vector of both batch structures. However, I am trying to use the segment_csr method to do my graph convolutions. This method requires to have the index sorted. Then what I do, is sort the batch vector and then with the output indexes, I sort the rest of the information of the new batch structure. I would like to know if you know a better way to combine the results of the both previous branches in a ordered way to be able to use the segment_csr. I didn't mentioned before, but the connections of the graph are generated using a knn approach.

[rusty1s]

Yes, if you want to use segment_csr, you will need to sort the matrices, and compute a new ptr tensor based on the sorted batch vector:

arange = torch.arange(int(batch.max()) + 1, device=batch.device)
ptr = torch.bucketize(arange, batch)

[rusty1s]

I think this is the most simple way. There exists alternative ways in which you could omit the sort by directly creating a N + M matrix that is sorted based on batches, in which you write values of both matrices into a single matrix according to the union of both batch vectors. That's tricky to implement though.

[dhorka]

Do you think that you could give me some hints about this method that you mention? I would like to give it a try. Maybe if I achieve that it could be a good method to have it as a utils in pytorch geometric.

[rusty1s]

Here is a small example:

batch1 = [0, 0, 0, 1, 1, 2, 2, 2]
batch2 = [0, 0, 1, 1, 2]

out = torch.empty(batch1.numel() + batch2.numel(), num_features)

index1 = [0, 1, 2, 5, 6, 9, 10, 11]
index2 = [3, 4, 7, 8, 12]

out[index1] = x1
out[index2] = x2

Here, you see that index1 and index2 hold the indices that dictate in which orders node features should be written to out. Those indices are based on the cumulated indices of both node feature matrices, e.g., values 0 in batch2 will be incremented by 3 (since there exists 3 zero values in batch1), and values 1 in batch1 will be incremented by 5 (since there exists 5 zero values in batch1 and batch2).

[dhorka]

I didn't see how to generate these index1 and index2 efficiently. For the moment, what I did is generate a mask in the dataset class that helps me to know in which part I should write. Instead of returning two data structures, my dataset class is returning 3 data structures, and the third one contains the mask to know in which position I should write.

If you have time I would appreciate a hint of how to generate these index1 and index 2 efficently just to improve my coding skills. Thanks!!!

[rusty1s]

import torch
from torch_geometric.utils import degree

batch1 = torch.tensor([0, 0, 0, 1, 1, 2, 2, 2])
batch2 = torch.tensor([0, 0, 1, 1, 2])

deg1 = degree(batch1, dtype=torch.long)
deg2 = degree(batch2, dtype=torch.long)
cumsum1 = torch.cat([torch.tensor([0]), deg1.cumsum(0)])
cumsum2 = torch.cat([torch.tensor([0]), deg2.cumsum(0)])

index1 = cumsum2[batch1] + torch.arange(batch1.numel())
print(index1)

index2 = cumsum1[1:][batch2] + torch.arange(batch2.numel())
print(index2)