gigl.src.common.models.pyg.link_prediction

Classes

LinkPredictionGNN

Link Prediction GNN model for both homogeneous and heterogeneous use cases Args: encoder (nn.Module): Either BasicGNN or Heterogeneous GNN for generating embeddings decoder (nn.Module): LinkPredictionDecoder for transforming embeddings into scores tasks (NodeAnchorBasedLinkPredictionTasks): Learning tasks for training (i.e. Retrieval, Margin, SSL, ...).