torch_geometric.nn.models.LightGCN

class LightGCN(num_nodes: int, embedding_dim: int, num_layers: int, alpha: Optional[Union[float, Tensor]] = None, **kwargs)[source]

Bases: Module

The LightGCN model from the “LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation” paper.

LightGCN learns embeddings by linearly propagating them on the underlying graph, and uses the weighted sum of the embeddings learned at all layers as the final embedding

\[\textbf{x}_i = \sum_{l=0}^{L} \alpha_l \textbf{x}^{(l)}_i,\]

where each layer’s embedding is computed as

\[\mathbf{x}^{(l+1)}_i = \sum_{j \in \mathcal{N}(i)} \frac{1}{\sqrt{\deg(i)\deg(j)}}\mathbf{x}^{(l)}_j.\]

Two prediction heads and training objectives are provided: link prediction (via link_pred_loss() and predict_link()) and recommendation (via recommendation_loss() and recommend()).

Note

Embeddings are propagated according to the graph connectivity specified by edge_index while rankings or link probabilities are computed according to the edges specified by edge_label_index.

Note

For an example of using LightGCN, see examples/lightgcn.py.

Parameters:

  • num_nodes (int) – The number of nodes in the graph.

  • embedding_dim (int) – The dimensionality of node embeddings.

  • num_layers (int) – The number of LGConv layers.

  • alpha (float or torch.Tensor, optional) – The scalar or vector specifying the re-weighting coefficients for aggregating the final embedding. If set to None, the uniform initialization of 1 / (num_layers + 1) is used. (default: None)

  • **kwargs (optional) – Additional arguments of the underlying LGConv layers.

forward(edge_index: Union[Tensor, SparseTensor], edge_label_index: Optional[Tensor] = None, edge_weight: Optional[Tensor] = None) → Tensor[source]

Computes rankings for pairs of nodes.

Parameters:

  • edge_index (torch.Tensor or SparseTensor) – Edge tensor specifying the connectivity of the graph.

  • edge_label_index (torch.Tensor, optional) – Edge tensor specifying the node pairs for which to compute rankings or probabilities. If edge_label_index is set to None, all edges in edge_index will be used instead. (default: None)

  • edge_weight (torch.Tensor, optional) – The weight of each edge in edge_index. (default: None)

reset_parameters()[source]

Resets all learnable parameters of the module.

get_embedding(edge_index: Union[Tensor, SparseTensor], edge_weight: Optional[Tensor] = None) → Tensor[source]

Returns the embedding of nodes in the graph.

Predict links between nodes specified in edge_label_index.

Parameters:

  • edge_index (torch.Tensor or SparseTensor) – Edge tensor specifying the connectivity of the graph.

  • edge_label_index (torch.Tensor, optional) – Edge tensor specifying the node pairs for which to compute probabilities. If edge_label_index is set to None, all edges in edge_index will be used instead. (default: None)

  • edge_weight (torch.Tensor, optional) – The weight of each edge in edge_index. (default: None)

  • prob (bool, optional) – Whether probabilities should be returned. (default: False)

recommend(edge_index: Union[Tensor, SparseTensor], edge_weight: Optional[Tensor] = None, src_index: Optional[Tensor] = None, dst_index: Optional[Tensor] = None, k: int = 1, sorted: bool = True) → Tensor[source]

Get top-\(k\) recommendations for nodes in src_index.

Parameters:

  • edge_index (torch.Tensor or SparseTensor) – Edge tensor specifying the connectivity of the graph.

  • edge_weight (torch.Tensor, optional) – The weight of each edge in edge_index. (default: None)

  • src_index (torch.Tensor, optional) – Node indices for which recommendations should be generated. If set to None, all nodes will be used. (default: None)

  • dst_index (torch.Tensor, optional) – Node indices which represent the possible recommendation choices. If set to None, all nodes will be used. (default: None)

  • k (int, optional) – Number of recommendations. (default: 1)

  • sorted (bool, optional) – Whether to sort the recommendations by score. (default: True)

Computes the model loss for a link prediction objective via the torch.nn.BCEWithLogitsLoss.

Parameters:
recommendation_loss(pos_edge_rank: Tensor, neg_edge_rank: Tensor, node_id: Optional[Tensor] = None, lambda_reg: float = 0.0001, **kwargs) → Tensor[source]

Computes the model loss for a ranking objective via the Bayesian Personalized Ranking (BPR) loss.

Note

The i-th entry in the pos_edge_rank vector and i-th entry in the neg_edge_rank entry must correspond to ranks of positive and negative edges of the same entity (e.g., user).

Parameters:

  • pos_edge_rank (torch.Tensor) – Positive edge rankings.

  • neg_edge_rank (torch.Tensor) – Negative edge rankings.

  • node_id (torch.Tensor) – The indices of the nodes involved for deriving a prediction for both positive and negative edges. If set to None, all nodes will be used.

  • lambda_reg (int, optional) – The \(L_2\) regularization strength of the Bayesian Personalized Ranking (BPR) loss. (default: 1e-4)

  • **kwargs (optional) – Additional arguments of the underlying torch_geometric.nn.models.lightgcn.BPRLoss loss function.