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()
andpredict_link()
) and recommendation (viarecommendation_loss()
andrecommend()
).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 byedge_label_index
.- Parameters
num_nodes (int) – The number of nodes in the graph.
embedding_dim (int) – The dimensionality of node embeddings.
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 of1 / (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) 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 toNone
, all edges inedge_index
will be used instead. (default:None
)
- get_embedding(edge_index: Union[Tensor, SparseTensor]) Tensor [source]
Returns the embedding of nodes in the graph.
- predict_link(edge_index: Union[Tensor, SparseTensor], edge_label_index: Optional[Tensor] = None, prob: bool = False) Tensor [source]
Predict links between nodes specified in
edge_label_index
.
- recommend(edge_index: Union[Tensor, SparseTensor], src_index: Optional[Tensor] = None, dst_index: Optional[Tensor] = None, k: int = 1) Tensor [source]
Get top-\(k\) recommendations for nodes in
src_index
.- Parameters
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
)
- link_pred_loss(pred: Tensor, edge_label: Tensor, **kwargs) Tensor [source]
Computes the model loss for a link prediction objective via the
torch.nn.BCEWithLogitsLoss
.- Parameters
pred (torch.Tensor) – The predictions.
edge_label (torch.Tensor) – The ground-truth edge labels.
**kwargs (optional) – Additional arguments of the underlying
torch.nn.BCEWithLogitsLoss
loss function.
- recommendation_loss(pos_edge_rank: Tensor, neg_edge_rank: Tensor, 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 theneg_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.
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.