Link prediction – recommending new friendships

In this section, we’ll implement a GAE for link prediction. This task aims to predict potential new connections in the network, which can be used to recommend new friendships among students.

First, let’s define our GAE model:

class LinkPredictor(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels):
        super().__init__()
        self.encoder = GCNConv(in_channels, hidden_channels)
    def encode(self, x, edge_index):
        return self.encoder(x, edge_index).relu()
    def decode(self, z, edge_label_index):
        return (
            z[edge_label_index[0]] * z[edge_label_index[1]]
     ...