class RGCNLayer(num_relations, input_dim=32, output_dim=None, use_bias=True, activation=None, activation_kwargs=None, self_loop_dropout=0.2, decomposition=None, decomposition_kwargs=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.

  • input_dim (int) – >0 the input dimension

  • num_relations (int) – the number of relations

  • output_dim (Optional[int]) – >0 the output dimension. If none is given, use the input dimension.

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

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

  • activation_kwargs (Optional[Mapping[str, Any]]) – 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 (Union[str, Decomposition, None]) – the decomposition to use, cf. Decomposition and decomposition_resolver

  • decomposition_kwargs (Optional[Mapping[str, Any]]) – 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, source, target, edge_type, edge_weights=None)[source]

Calculate enriched entity representations.

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

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

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

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

  • edge_weights (Optional[FloatTensor]) – shape: (num_triples,) Scalar edge weights per triple.

Return type:



shape: (num_entities, output_dim) Enriched entity representations.