How to use disjoint_train_ratio properly in RandomLinkSplit?

[nabihach]

I'm doing link prediction over the OGB_MAG dataset. I split the dataset into train/val/test with the RandomLinkSplit transform, but I am a bit unsure how to use the disjoint_train_ratio variable properly. To be more concrete, I know that during training we should use the message edges to predict the supervision edges, but how exactly do I do that? I am pasting the script I have so far. Would be great if someone can point out how to actually train over message edges only and not supervision edges. Anything else that stands out as problematic or erroneous, please feel free to point that out as well. Thanks!

import torch
import os
from torch_geometric.datasets import OGB_MAG
from torch_geometric.data import Data
from torch_geometric.loader import NeighborLoader
from torch_geometric.nn import GCNConv
import torch_geometric.transforms as T
import torch.nn.functional as F
from tqdm import tqdm

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) #x.log_softmax(dim=-1)

# Load the dataset

path = os.path.join('..', 'data', 'OGB_MAG')
mag_dataset = OGB_MAG(path, preprocess='metapath2vec')
mag_data = mag_dataset[0]

# Add reverse edges to the heterogeneous graph.

r, c = mag_data.edge_index_dict[('author', 'affiliated_with', 'institution')]
mag_data[('institution', 'to', 'author')].edge_index = torch.stack([c, r])
r, c = mag_data.edge_index_dict[('author', 'writes', 'paper')]
mag_data[('paper', 'to', 'author')].edge_index = torch.stack([c, r])
r, c = mag_data.edge_index_dict[('paper', 'has_topic', 'field_of_study')]
mag_data[('field_of_study', 'to', 'paper')].edge_index = torch.stack([c, r])
r, c = mag_data.edge_index_dict[('paper', 'cites', 'paper')]
mag_data[('paper', 'revcites', 'paper')].edge_index = torch.stack([c, r])

# Split data intro train/val/test

tfs = T.RandomLinkSplit(is_undirected=False, 
                      add_negative_train_samples=True,
                      disjoint_train_ratio=0.2,
                      split_labels=False,
                      edge_types=[('author', 'affiliated_with', 'institution'), ('author', 'writes', 'paper'), ('paper', 'cites', 'paper'), ('paper', 'has_topic', 'field_of_study')], 
                      rev_edge_types=[('institution', 'to', 'author'), ('paper', 'to', 'author'), ('paper', 'revcites', 'paper'), ('field_of_study', 'to', 'paper')], 
                      )
train_data, val_data, test_data = tfs(mag_data)

# Create training data

train_x= torch.cat(  (train_data['paper'].x, train_data['author'].x, train_data['field_of_study'].x, train_data['institution'].x), 0).float()
train_y = torch.cat( (train_data[('paper', 'cites', 'paper')].edge_label, train_data[('author', 'writes', 'paper')].edge_label, train_data[('paper', 'has_topic', 'field_of_study')].edge_label, train_data[('author', 'affiliated_with', 'institution')].edge_label),  0)
train_edge_index = torch.cat( (train_data[('paper', 'cites', 'paper')].edge_label_index, train_data[('author', 'writes', 'paper')].edge_label_index, train_data[('paper', 'has_topic', 'field_of_study')].edge_label_index, train_data[('author', 'affiliated_with', 'institution')].edge_label_index),  1)
training_data = Data(x=train_x, y=train_y, edge_index=train_edge_index)
training_data = T.ToSparseTensor(remove_edge_index=False)(training_data)
training_data

# Define model and hyperparams

hidden_channels = 128
num_layers = 1
dropout = 0.5
device = 'cuda'

model = GCNConv(hidden_channels, hidden_channels, cached=False).to(device)
predictor = LinkPredictor(hidden_channels, hidden_channels, 1, num_layers, dropout).to(device)
optimizer = torch.optim.Adam(list(model.parameters()) + list(predictor.parameters()), lr=0.01)

# Define batch loader

input_nodes = torch.ones(train_data.num_nodes, dtype=torch.bool)
kwargs = {'batch_size': 256, 'num_workers': 6, 'persistent_workers': True}
train_loader = NeighborLoader(training_data, num_neighbors=[20], shuffle=True, directed=True, input_nodes=input_nodes, **kwargs)

def train(model, predictor, data, optimizer):
    model.train()
    predictor.train()

    total_loss = total_examples = 0
    for batch in tqdm(train_loader):
        optimizer.zero_grad()
        
        batch = batch.to(device)
        batch_size = batch.batch_size
        h = model(batch.x, batch.adj_t)
        
        edges = batch.edge_index
        labels = batch.y
        out = predictor(h[edges[0]], h[edges[1]])

        loss = F.binary_cross_entropy(out.squeeze(-1), labels)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        torch.nn.utils.clip_grad_norm_(predictor.parameters(), 1.0)
        optimizer.step()

        num_examples = out.size(0)
        total_loss += loss.item() * num_examples
        total_examples += num_examples

    return total_loss / total_examples

# Main loop 

model.reset_parameters()
predictor.reset_parameters()
for epoch in range(0, 6):
    loss = train(model, predictor, training_data, optimizer)
    print("Epoch: "+ str(epoch) + ", Loss: " + str(loss))

[rusty1s]

Thanks for your interest. Can you take a look at #4026 and see if that resolves your issues? We are currently working on better link-prediction support within NeighborLoader.

[nabihach]

Thanks for your quick response @rusty1s . I know I've raised a lot of questions about NeighborLoader lately :) This particular question is not about that. Here I am inquiring about the RandomLinkSplit transform -- it has an argument disjoint_train_ratio where we specify the ratio of message edges vs supervision edges:

tfs = T.RandomLinkSplit(is_undirected=False, 
                      add_negative_train_samples=True,
                      disjoint_train_ratio=0.2,
                      split_labels=False,
                      edge_types=[('author', 'affiliated_with', 'institution'), ('author', 'writes', 'paper'), ('paper', 'cites', 'paper'), ('paper', 'has_topic', 'field_of_study')], 
                      rev_edge_types=[('institution', 'to', 'author'), ('paper', 'to', 'author'), ('paper', 'revcites', 'paper'), ('field_of_study', 'to', 'paper')], 
                      )
train_data, val_data, test_data = tfs(mag_data)

Here I specify disjoint_train_ratio to be a non-zero value, so my expectation was that tfs would return training_message_data, training_supervision_data, val_data and test_data but it does not. How do I use train_data here to distinguish message edges from supervision edges?

[rusty1s]

Ah I see. The distinction between message passing edges and supervision edges happens in train_data.edge_index (message passing) and train_data.edge_label_index (supervision). In case djsjoint_train_ratio is not set, both will be identical.