Hetero Link Prediction: retrieve final predicted edge_index

[wgeul]

Hi,

In the homogeneous GAE example (https://github.com/pyg-team/pytorch_geometric/blob/master/examples/link_pred.py) the last two lines show how to retrieve the 'fitted' edge_index as follows:

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

Question: How do I obtain the fitted edge index in the heterogeneous example?

I notice that in the heterogeneous example (https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/hetero_link_pred.py) the objective is slightly different as it's fitting to the given ratings specified as an edge_label. However, I'm actually interested in the same objective as in the homogeneous variant: I'd like to predict the edge_index itself by optimizing the network to fit the labelled probabilities of positive/negative edge links. This requires some sort of transformation of the encoded z_dict to an edge_index (see below function decode_all).

I've made the following alterations and would like to ask your advice particularly on the decode_all function within the Model class and I wonder whether the forward function in EdgeEncoder should also be adjusted to reflect binary cross entropy objective?

import os.path as osp

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

import torch_geometric.transforms as T
from torch_geometric.datasets import MovieLens
from torch_geometric.utils import negative_sampling
from torch_geometric.nn import SAGEConv, to_hetero


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

__file__ = r'<local dir>'
url = 'https://files.grouplens.org/datasets/movielens/ml-latest-small.zip'
path = osp.join(osp.dirname(osp.realpath(__file__)), 'MovieLens')
dataset = MovieLens(path, model_name='all-MiniLM-L6-v2')
data = dataset[0].to(device)

# Add user node features for message passing:
data['user'].x = torch.eye(data['user'].num_nodes, device=device)
del data['user'].num_nodes

##Remove edge_label
del data['user', 'rates', 'movie'].edge_label

# Add a reverse ('movie', 'rev_rates', 'user') relation for message passing:
data = T.ToUndirected()(data)
data = T.ToDevice(device)(data)

# Perform a link-level split into training, validation, and test edges:
train_data, val_data, test_data =  T.RandomLinkSplit(
        num_val=0.05,
        num_test=0.1,
        is_undirected=True,
        add_negative_train_samples=True,
        edge_types=[('user', 'rates', 'movie')],
        rev_edge_types=[('movie', 'rev_rates', 'user')])(data)


target_edge = ['user', 'movie']


class GNNEncoder(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

class EdgeDecoder(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.lin1 = Linear(2 * hidden_channels, hidden_channels)
        self.lin2 = Linear(hidden_channels, 1)

    def forward(self, z_dict, edge_label_index):
        row, col = edge_label_index
        z = torch.cat([z_dict[target_edge[0]][row], z_dict[target_edge[1]][col]], dim=-1)

        z = self.lin1(z).relu()
        z = self.lin2(z)
        return z.view(-1)



class Model(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.encoder = GNNEncoder(hidden_channels, hidden_channels)
        self.encoder = to_hetero(self.encoder, data.metadata(), aggr='sum')
        self.decoder = EdgeDecoder(hidden_channels)

    def forward(self, x_dict, edge_index_dict, edge_label_index):
        z_dict = self.encoder(x_dict, edge_index_dict)
        return self.decoder(z_dict, edge_label_index)
    
    def decode_all(self, x_dict, edge_index_dict, edge_label_index):
        z = self.encoder(x_dict, edge_index_dict)
        #???
        #z is dict, how to deal with this?
        #function taken from link_pred.py
        mm = []
        for k in z:
            mm.append(z[k] @ z[k].t())     
        prob_adj = z @ z.t()
        return (prob_adj > 0).nonzero(as_tuple=False).t()

model = Model(hidden_channels=32).to(device)

# Due to lazy initialization, we need to run one model step so the number
# of parameters can be inferred:
with torch.no_grad():
    model.encoder(train_data.x_dict, train_data.edge_index_dict)

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

def train():
    model.train()
    optimizer.zero_grad()

    pred = model(train_data.x_dict, train_data.edge_index_dict,
                 train_data[target_edge].edge_label_index)
    target = train_data[target_edge].edge_label
    loss = torch.nn.BCEWithLogitsLoss()(pred, target)
    loss.backward()
    optimizer.step()
    return float(loss)


@torch.no_grad()
def test(data):
    model.eval()
    pred = model(data.x_dict, data.edge_index_dict,
                 data[target_edge].edge_label_index)
    pred = pred.clamp(min=0, max=1)
    target = data[target_edge].edge_label
    loss = torch.nn.BCEWithLogitsLoss()(pred, target)
    return float(loss)


for epoch in range(1, 11):
    loss = train()
    train_rmse = test(train_data)
    val_rmse = test(val_data)
    test_rmse = test(test_data)
    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Train: {train_rmse:.4f}, '
          f'Val: {val_rmse:.4f}, Test: {test_rmse:.4f}')

final_edge_index = model.decode_all(test_data.x_dict, test_data.edge_index_dict,
                 test_data[target_edge].edge_label_index)

[rusty1s]

I think your general procedure looks correct. You would either iterate over each pair of node type and compute/decode the local adjacencies between node types. Alternatively, you can concatenate oll matrices inside your dict in the first dimension, and then just do reconstruction as in the homogeneous case.

[rusty1s]

Link prediction refers to the problem of adding missing edges. If you want to infer a graph from no edges, one usually starts with a fully-connected graph and tries to learn to sparsify it. As such, I think your approach is correct :)

[wgeul]

Hi @rusty1s, thanks for your answer and the advice on the approach to link prediction with missing edges.

[jameswpm]

Link prediction refers to the problem of adding missing edges. If you want to infer a graph from no edges, one usually starts with a fully-connected graph and tries to learn to sparsify it. As such, I think your approach is correct :)

Hello @rusty1s

I don't know if reactivating old discussions is a good approach, but let me know if creating a new question is better.

I'm highly interested in implementing the aforementioned problem (training a heterogeneous GNN with a fully connected graph and then predicting links in a disconnected one). Still, I found no good example/tutorial/paper about it. Most of them are in the line of recommendation of non-existent links in the initial graph.

Do you have any source code or material to deal with this kind of problem? Your explanation in the quoted message is a little bit broad for me (sorry for the noob question, I started with PyG some months ago, and I'm still lacking in some terminology and basic stuff)

[rusty1s]

One paper related to this is https://arxiv.org/pdf/1802.04687.pdf. Ofc, there exists other approaches to do this as well, e.g., https://openreview.net/pdf?id=SJxYOVSgUB. I don't have PyG source code for these though :(

[jameswpm]

Thanks a lot for the links. I lacked the term "Graph reconstruction," in my first searches but I've already clarified the theme just with the papers' abstracts. I think it will help a lot.