How to predict direction of edge in a graph?

[urmi22]

Hi,

I am trying to classify if there is an edge between A -> B or B -> A in a graph. However, I am not very sure how do I go about designing the model and the graph. Can you anyone please suggest some plausible route I can take?

Thanks in advance!

[rusty1s]

For undirected graphs, this is mostly covered in our link prediction example, see here. To convert this into link prediction example for directed graphs, one needs to replace the decoder to account for that, e.g., via:

def decode(self, z, edge_label_index):
        return self.MLP(torch.cat([z[edge_label_index[0], z[edge_label_index[1]], dim=-1))

Hope this helps :)

[AdarshMJ]

Hi @rusty1s @EdisonLeeeee

Thank you for this answer. I also have a similar query. I did what you suggested. I have directed graphs and they look like this -

(Data(x=[176, 6], edge_index=[2, 400], edge_attr=[400, 1], y=[176], edge_label=[400], edge_label_index=[2, 400]), Data(x=[176, 6], edge_index=[2, 400], edge_attr=[400, 1], y=[176], edge_label=[46], edge_label_index=[2, 46]), Data(x=[176, 6], edge_index=[2, 423], edge_attr=[423, 1], y=[176], edge_label=[94], edge_label_index=[2, 94]))

To account for the directions I did -

from torch.nn import Linear
import torch.nn.functional as F

class Net(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)
        self.MLP = Linear(2*out_channels, 1)
    
    def encode(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        return self.conv2(x, edge_index)

    def decode(self, z, edge_label_index):
        #return self.MLP(torch.cat([z[edge_label_index[0], z[edge_label_index[1]]]], dim=-1))
        return self.MLP(torch.cat((z[edge_label_index[0]],z[edge_label_index[1]]), dim=-1))

    def decode_all(self, z):
        prob_adj = z @ z.t()
        return (prob_adj > 0).nonzero(as_tuple=False).t()


model = Net(6, 128, 128).to(device)
print(model)
optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)
criterion = torch.nn.BCEWithLogitsLoss()

I only have total 19 graphs split amongst train,test and val. I am getting accuracy such as this -

Epoch: 096, Loss: 0.3753, Val: 0.9112, Test: 0.8821
Epoch: 097, Loss: 0.3717, Val: 0.9121, Test: 0.8816
Epoch: 098, Loss: 0.3917, Val: 0.9127, Test: 0.8819
Epoch: 099, Loss: 0.3870, Val: 0.9127, Test: 0.8826
Epoch: 100, Loss: 0.3807, Val: 0.9127, Test: 0.8831
Final Test: 0.8824

My question is how do I use this trained model to perform inference? Like if I have a new data, say a graph with just two nodes and I want to predict if a link exists between them and which direction. How would I go about doing that? I dont fully understand what this does. Thank you for your time!

z = model.encode(test_data.x, test_data.edge_index)
final_edge_index = model.decode_all(z)
print(final_edge_index)

tensor([[  0,   0,   0,  ..., 175, 175, 175],
        [  0,   1,   2,  ..., 173, 174, 175]])

[EdisonLeeeee]

You might need to change the code in decode_all to predict the "edge probability":

    def decode_all(self, z):
        prob_adj = z @ z.t()
        return (prob_adj.sigmoid() > 0.5).nonzero(as_tuple=False).t()

[AdarshMJ]

Thank you for the reply. But even after changing the line the output for this is the same -

z = model.encode(test_data.x, test_data.edge_index)
final_edge_index = model.decode_all(z)
print(final_edge_index)

tensor([[  0,   0,   0,  ..., 175, 175, 175],
        [  0,   1,   2,  ..., 173, 174, 175]])

Could you please explain what is this referring to or what does this mean? Do the zeros indicate no edge?

[EdisonLeeeee]

Could you please explain what is this referring to or what does this mean?

Since the model is optimized with BCEWithLogitsLoss that uses sigmoid activation, the output during inference should come with sigmoid as well. BTW, I feel you should use the same decoding scheme for training and inference - decode_all uses dot-product but decode uses MLP+concat.

Do the zeros indicate no edge?

It depends on how the model was trained with labels. IMO, outputs greater than 0.5 indicate positive edges and negative ones otherwise.