Link Prediction task training very slow on OGB-Citation2, is it because using linkneighborloader?

[Xurorx]

Hello PyG team @rusty1s ,

I am implementing a link prediction task using GCN or GraphSage on ogbl-citation2 dataset, but it takes 12h to run just 1 epoch. Is it because I am using LinkNeighborLoader? I took a reference of PyG teams' code. I thought my implementation is pretty similar to it, but why training takes such long time?

Device I am using:
NVIDIA Tesla V100 NVLINK, 32 GB of GPU memory

My code is as follows,

import os
import torch
import torch_geometric
from torch_geometric.loader import LinkNeighborLoader
from torch.utils.data import DataLoader
import torch.nn.functional as F
from ogb.linkproppred import PygLinkPropPredDataset, Evaluator
from torch_geometric.utils import to_undirected

# import from local package
from model import *
from argument.aml_adapter import get_args_from_module_spec
from logger import get_logger, tracing_exception
from utils import EarlyStopping


logger = get_logger(__name__)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logger.info(f"Device: '{device}'")

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

    total_loss = total_examples = 0
    for sampled_data in train_loader:
        sampled_data.to(device)
        optimizer.zero_grad()
        
        emb = model(sampled_data.x, sampled_data.edge_index)
        src, dst = sampled_data.edge_index
        pos_out = predictor(emb[src], emb[dst])
        pos_loss = -torch.log(pos_out + 1e-15).mean()

        # Just do some trivial random sampling.
        dst_neg = torch.randint(0, sampled_data.x.size(0), src.size(),
                                dtype=torch.long, device=device)
        neg_out = predictor(emb[src], emb[dst_neg])
        neg_loss = -torch.log(1 - neg_out + 1e-15).mean()
        loss = pos_loss + neg_loss

        loss.backward()
        optimizer.step()

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

    return total_loss / total_examples

@torch.no_grad()
def test(model, predictor, data, split_edge, evaluator, batch_size, split):
    model.eval()
    predictor.eval()
    model = model.to('cpu')
    predictor = predictor.to('cpu')
    
    emb = model(data.x, data.edge_index)

    source = split_edge[split]['source_node']
    target = split_edge[split]['target_node']
    target_neg = split_edge[split]['target_node_neg']

    pos_preds = []
    for perm in DataLoader(range(source.size(0)), batch_size):
        src, dst = source[perm], target[perm]
        pos_preds += [predictor(emb[src], emb[dst]).squeeze().cpu()]
    pos_pred = torch.cat(pos_preds, dim=0)

    neg_preds = []
    source = source.view(-1, 1).repeat(1, 1000).view(-1)
    target_neg = target_neg.view(-1)
    for perm in DataLoader(range(source.size(0)), batch_size):
        src, dst_neg = source[perm], target_neg[perm]
        neg_preds += [predictor(emb[src], emb[dst_neg]).squeeze().cpu()]
    neg_pred = torch.cat(neg_preds, dim=0).view(-1, 1000)

    return evaluator.eval({
        'y_pred_pos': pos_pred,
        'y_pred_neg': neg_pred,
    })['mrr_list'].mean().item()

@tracing_exception(logger=logger)
def main():
    """ script main function """
    logger.info("Starting ...")
    logger.info("Reading parameters...")
    args = get_args_from_module_spec(module_spec_file_name='module_spec.yaml')
    logger.info("Finished reading parameters.")
    os.makedirs(args.output_path, exist_ok=True)

    # Download and process data
    dataset = PygLinkPropPredDataset(name="ogbl-citation2", root=args.input_path)
    split_edge = dataset.get_edge_split()
    data = dataset[0]

    data.edge_index = to_undirected(data.edge_index, data.num_nodes)

    train_loader = LinkNeighborLoader(
        data=data,
        num_neighbors=[15,10,5],
        edge_label_index=data.edge_index,
        batch_size=args.batch_size
    )

    # Initialize model
    if args.model == 'Sage':
        model = SAGE(in_channels=data.num_features, hidden_channels=args.hidden_dim, out_channels=args.hidden_dim, num_layers=args.num_layers, dropout=args.dropout).to(device)
    elif args.model =="GCN":
        model = GCN(in_channels=data.num_features, hidden_channels=args.hidden_dim, out_channels=args.hidden_dim, num_layers=args.num_layers, dropout=args.dropout).to(device)

    logger.info(f"model: {model}")
    predictor = LinkPredictor(in_channels=args.hidden_dim, hidden_channels=args.hidden_dim, out_channels=1, num_layers=args.num_layers, dropout=args.dropout).to(device)
    evaluator = Evaluator(name='ogbl-citation2')

    ### Training a Homogeneous Link-level GNN ### 

    if not args.checkpoint_path:
        optimizer = torch.optim.Adam(list(model.parameters()) + list(predictor.parameters()), lr=args.learning_rate)
        early_stopping = EarlyStopping(patience=7, metric="mrr", verbose=True, path=os.path.join(args.output_path,"checkpoint.pt"), logger=logger.info)
        num_epochs = args.epochs + 1

        logger.info(f"Starting training.")
        for epoch in range(1, num_epochs):

            # Train
            loss = train(model, predictor, train_loader, optimizer)
            logger.info(f"Epoch: {epoch:03d}, Loss: {loss:.4f}")

            if epochs % 10 == 0:
                # Validation MRR
                val_result = test(model, predictor, data, split_edge, evaluator, 12 * 1024, 'valid')

            # # Early stopping on MRR
            # state_dict = {'model_state_dict': model.state_dict(), 'predictor_state_dict': predictor.state_dict()}
            # early_stopping(val_result, state_dict)
            # if early_stopping.early_stop:
            #     logger.info("Early stopping")
            #     break

    else:
        # Load the checkpoint
        checkpoint_path = os.path.join(args.checkpoint_path, "checkpoint.pt")
        logger.info(f"Load checkpoint from: {checkpoint_path}.")
        checkpoint = torch.load(checkpoint_path)
        model.load_state_dict(checkpoint['model_state_dict'])
        predictor.load_state_dict(checkpoint['predictor_state_dict'])
        model = model.to(device)
        predictor = predictor.to(device)

    # Evaluation
    logger.info(f"Start evaluation.")
    test_result = test(model, predictor, data, split_edge, evaluator, 12 * 1024, 'test')
    logger.info(f"Test MRR: {test_result:.4f}")


if __name__ == "__main__":
    main()

This is model.py

import torch
from torch_geometric.nn import GCNConv, to_hetero, SAGEConv
import torch.nn.functional as F

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

        self.convs = torch.nn.ModuleList()
        if num_layers == 1:
            self.convs.append(
                SAGEConv(in_channels, out_channels, normalize=False))
        else:
            self.convs.append(SAGEConv(in_channels, hidden_channels))
            for _ in range(num_layers - 2):
                self.convs.append(SAGEConv(hidden_channels, hidden_channels))
            self.convs.append(SAGEConv(hidden_channels, out_channels))

        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 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()
        if num_layers == 1:
            self.convs.append(
                GCNConv(in_channels, out_channels, normalize=False))
        else:
            self.convs.append(
                GCNConv(in_channels, hidden_channels, normalize=False))
            for _ in range(num_layers - 2):
                self.convs.append(
                    GCNConv(hidden_channels, hidden_channels, normalize=False))
            self.convs.append(
                GCNConv(hidden_channels, out_channels, normalize=False))

        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()
        if num_layers == 1:
            self.lins.append(torch.nn.Linear(in_channels, out_channels))
        else:
            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, source_emb, dest_emb):
        x = source_emb * dest_emb
        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)

And the params that I am using are:

epochs: 150
batch_size: 4096
learning_rate: 0.0008
model: Sage
num_layers: 2
hidden_dim: 128
dropout: 0.0

[rusty1s]

Are you using pyg-lib? Also, you want to increase the number of workers in LinkNeighborLoader.

[Xurorx]

Hi @rusty1s I am not using pyg-lib, just torch_geometric. Let me try increasing number of worker in LinkNeighborLoader. Also I wonder if setting num_workers =1 is different from not setting num_workers? I did some research, looks like if not set num_workers, its default value is 0. And num_workers = 1 will turn on multi-process while 0 won't? Just want to confim.

[rusty1s]

Yeah, you should definitely try out pyg-lib. For num_workers=1, it will use multi-processing (but on a single thread), so num_workers=1 usually decreases your throughput. Try something like num_workers=8 at best.

[Xurorx]

Thank you @rusty1s ! Increaseing num_workers help reduce the training time a lot. I will try pyg-lib. I haven't used pyg-lib before, is pyg-lib more efficient?

[rusty1s]

Yes, it introduces some major speed-ups to the sampling routine :)

[Xurorx]

Thanks a lot!