Advice to Imporve the code for link prediction without node features

[bis1999]

Below is the code

The problem is that I have a training set containing edges and nodes without any node or edge features, I want to compare GNN with other models as well, so basically I take 0.80 of data and train the GCN on it, and test it on some examples of neg and pos examples. I want to have advice on how I can improve it, I have taken this code from the ogbl-ddi page written by @rusty1s https://github.com/snap-stanford/ogb/tree/master/examples/linkproppred/ddi, Let me also know if there some issues you see.

import argparse

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch_geometric.utils import negative_sampling

import torch_geometric.transforms as T
from torch_geometric.nn import GCNConv, SAGEConv

from ogb.linkproppred import PygLinkPropPredDataset, Evaluator
from torch_sparse import SparseTensor


import pickle
import random







# Path to the file where the dictionary is saved
file_path = 'OLP_updated.pickle'

# Load the dictionary from the file using pickle
with open(file_path, 'rb') as file:
    new = pickle.load(file)

edges = new[new['network_index']== 17.0]["edges_id"].tolist()[0]
num_nodes = int(new[new['network_index']== 17.0]["number_nodes"].tolist()[0])

edges_pair = [(x[0],x[1]) for x in edges ]

sample_ = int(len(edges_pair)*0.20)
edges_train_n = len(edges_pair) - sample_

train_edges = random.sample(edges_pair, edges_train_n)
edges_not_in_train = list(set(edges_pair) - set(train_edges))


row_ = [ i[0] for i in train_edges]
cols_ = [ i[1] for i in train_edges]


sparse_tensor = SparseTensor(row=torch.tensor(row_),
                             col=torch.tensor(cols_), sparse_sizes=(num_nodes,int(num_nodes)))
adj_t = sparse_tensor

class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(GCN, self).__init__()

        self.convs = torch.nn.ModuleList()
        self.convs.append(GCNConv(in_channels, hidden_channels, cached=True))
        for _ in range(num_layers - 2):
            self.convs.append(
                GCNConv(hidden_channels, hidden_channels, cached=True))
        self.convs.append(GCNConv(hidden_channels, out_channels, cached=True))

        self.dropout = dropout

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()

    def forward(self, x, adj_t):
        for conv in self.convs[:-1]:
            x = conv(x, adj_t)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.convs[-1](x, adj_t)
        return x
    
    
    
class LinkPredictor(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(LinkPredictor, self).__init__()

        self.lins = torch.nn.ModuleList()
        self.lins.append(torch.nn.Linear(in_channels, hidden_channels))
        for _ in range(num_layers - 2):
            self.lins.append(torch.nn.Linear(hidden_channels, hidden_channels))
        self.lins.append(torch.nn.Linear(hidden_channels, out_channels))

        self.dropout = dropout

    def reset_parameters(self):
        for lin in self.lins:
            lin.reset_parameters()

    def forward(self, x_i, x_j):
        x = x_i * x_j
        for lin in self.lins[:-1]:
            x = lin(x)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.lins[-1](x)
        return torch.sigmoid(x)

    
def train(model, predictor, x, adj_t, optimizer):

    row, col, _ = adj_t.coo()
    edge_index = torch.stack([col, row], dim=0)

    model.train()
    predictor.train()


    total_loss = total_examples = 0
    
    optimizer.zero_grad()

    h = model(x, adj_t)


    pos_out = predictor(h[edge_index[0]], h[edge_index[1]])
    pos_loss = -torch.log(pos_out + 1e-15).mean()

    edge = negative_sampling(edge_index, num_nodes=num_nodes,
                             num_neg_samples=edge_index.size(1), method='dense')

    neg_out = predictor(h[edge[0]], h[edge[1]])
    neg_loss = -torch.log(1 - neg_out + 1e-15).mean()

    loss = pos_loss + neg_loss
    loss.backward()

    torch.nn.utils.clip_grad_norm_(x, 1.0)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    torch.nn.utils.clip_grad_norm_(predictor.parameters(), 1.0)

    optimizer.step()

   


    return loss.item()




num_layers = 4
hidden_channels = 256
dropout = 0.01
lr = 0.001
device ="cpu"


model = GCN(hidden_channels, hidden_channels,hidden_channels, num_layers,dropout).to(device)
emb = torch.nn.Embedding(adj_t.size(0),hidden_channels).to(device)
predictor = LinkPredictor(hidden_channels, hidden_channels, 1,num_layers, dropout).to(device)


optimizer = torch.optim.Adam(
            list(model.parameters()) + list(emb.parameters()) +
            list(predictor.parameters()), lr=lr)


x = emb.weight


loss_graph = []

for i in range(1000):
    ll = train(model, predictor, x, adj_t, optimizer)
    loss_graph.append(ll)
    print(ll)
    
    
from sklearn.metrics import roc_auc_score
tt = new["df_tr_top"]
neg_ = tt[tt["Link_present"] == 0]
negatve_examples = torch.tensor(neg_["non_links"].tolist()).T
positive_examples = torch.tensor(edges_not_in_train).T
positive_examples.size(1)

test_edge_label_index = torch.cat(
        [positive_examples, negatve_examples],
        dim=-1,

    )

test_edge_label = torch.cat([
        torch.ones(positive_examples.size(1)),
        torch.zeros(negatve_examples.size(1))
    ], dim=0)

model.eval()
predictor.eval()


h = model(x, adj_t)
pred = predictor(h[test_edge_label_index[0]],h[test_edge_label_index[1]])
roc_auc_score(test_edge_label.numpy(),pred.view(-1).detach().numpy())

[EdisonLeeeee]

You can also consider exploring unsupervised methods (e.g., DeepWalk or Node2Vec) to generate node features first and then feed them into GNNs, FYI: https://github.com/pyg-team/pytorch_geometric/blob/master/examples/node2vec.py

[bis1999]

Thank you so that sounds great!