RGCNLayer

class RGCNLayer(num_relations: int, input_dim: int = 32, output_dim: int | None = None, use_bias: bool = True, activation: str | Module | None = None, activation_kwargs: Mapping[str, Any] | None = None, self_loop_dropout: float = 0.2, decomposition: str | Decomposition | None = None, decomposition_kwargs: Mapping[str, Any] | None = None)[source]

Bases: Module

An RGCN layer from [schlichtkrull2018] updated to match the official implementation.

This layer uses separate decompositions for forward and backward edges (i.e., “normal” and implicitly created inverse relations), as well as a separate transformation for self-loops.

Ignoring dropouts, decomposition and normalization, it can be written as

\[y_i = \sigma( W^s x_i + \sum_{(e_j, r, e_i) \in \mathcal{T}} W^f_r x_j + \sum_{(e_i, r, e_j) \in \mathcal{T}} W^b_r x_j + b )\]

where \(b, W^s, W^f_r, W^b_r\) are trainable weights. \(W^f_r, W^b_r\) are relation-specific, and commonly enmploy a weight-sharing mechanism, cf. Decomposition. \(\sigma\) is an activation function. The individual terms in both sums are typically weighted. This is implemented by EdgeWeighting. Moreover, RGCN employs an edge-dropout, however, this needs to be done outside of an individual layer, since the same edges are dropped across all layers. In contrast, the self-loop dropout is layer-specific.

Initialize the layer.

Parameters:

  • input_dim (int) – >0 the input dimension

  • num_relations (int) – the number of relations

  • output_dim (int | None) – >0 the output dimension. If none is given, use the input dimension.

  • use_bias (bool) – whether to use a trainable bias

  • activation (str | Module | None) – the activation function to use. Defaults to None, i.e., the identity function serves as activation.

  • activation_kwargs (Mapping[str, Any] | None) – additional keyword-based arguments passed to the activation function for instantiation

  • self_loop_dropout (float) – 0 <= self_loop_dropout <= 1 the dropout to use for self-loops

  • decomposition (str | Decomposition | None) – the decomposition to use, cf. Decomposition and decomposition_resolver

  • decomposition_kwargs (Mapping[str, Any] | None) – the keyword-based arguments passed to the decomposition for instantiation

Methods Summary

forward(x, source, target, edge_type[, ...])

Calculate enriched entity representations.

Methods Documentation

forward(x: Tensor, source: Tensor, target: Tensor, edge_type: Tensor, edge_weights: Tensor | None = None) → Tensor[source]

Calculate enriched entity representations.

Parameters:

  • x (Tensor) – shape: (num_entities, input_dim) The input entity representations.

  • source (Tensor) – shape: (num_triples,) The indices of the source entity per triple.

  • target (Tensor) – shape: (num_triples,) The indices of the target entity per triple.

  • edge_type (Tensor) – shape: (num_triples,) The relation type per triple.

  • edge_weights (Tensor | None) – shape: (num_triples,) Scalar edge weights per triple.

Returns:

shape: (num_entities, output_dim) Enriched entity representations.

Return type:

Tensor