Models

Implementations of various knowledge graph embedding models.

Name

Reference

Citation

ComplEx

pykeen.models.ComplEx

[trouillon2016]

ComplExLiteral

pykeen.models.ComplExLiteral

[agustinus2018]

ConvE

pykeen.models.ConvE

[dettmers2018]

ConvKB

pykeen.models.ConvKB

[nguyen2018]

DistMult

pykeen.models.DistMult

[yang2014]

DistMultLiteral

pykeen.models.DistMultLiteral

[agustinus2018]

ERMLP

pykeen.models.ERMLP

[dong2014]

ERMLPE

pykeen.models.ERMLPE

[sharifzadeh2019]

HolE

pykeen.models.HolE

[nickel2016]

KG2E

pykeen.models.KG2E

[he2015]

NTN

pykeen.models.NTN

[socher2013]

ProjE

pykeen.models.ProjE

[shi2017]

RESCAL

pykeen.models.RESCAL

[nickel2011]

RGCN

pykeen.models.RGCN

[schlichtkrull2018]

RotatE

pykeen.models.RotatE

[sun2019]

SimplE

pykeen.models.SimplE

[kazemi2018]

StructuredEmbedding

pykeen.models.StructuredEmbedding

[bordes2011]

TransD

pykeen.models.TransD

[ji2015]

TransE

pykeen.models.TransE

[bordes2013]

TransH

pykeen.models.TransH

[wang2014]

TransR

pykeen.models.TransR

[lin2015]

TuckER

pykeen.models.TuckER

[balazevic2019]

UnstructuredModel

pykeen.models.UnstructuredModel

[bordes2014]

Note

This table can be re-generated with pykeen ls models -f rst

class pykeen.models.ComplEx(triples_factory, embedding_dim=200, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of ComplEx [trouillon2016].

See also

Official implementation: https://github.com/ttrouill/complex/

Initialize the module.

Parameters
  • triples_factory (TriplesFactory) – TriplesFactory The triple factory connected to the model.

  • embedding_dim (int) – int The embedding dimensionality of the entity embeddings.

  • automatic_memory_optimization (Optional[bool]) – bool Whether to automatically optimize the sub-batch size during training and batch size during evaluation with regards to the hardware at hand.

  • loss (Optional[_Loss]) – OptionalLoss (optional) The loss to use. Defaults to SoftplusLoss.

  • preferred_device (Optional[str]) – str (optional) The default device where to model is located.

  • random_seed (Optional[int]) – int (optional) An optional random seed to set before the initialization of weights.

  • regularizer (Optional[Regularizer]) – BaseRegularizer The regularizer to use.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 300, 'low': 50, 'q': 50, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

static interaction_function(h, r, t)[source]

Evaluate the interaction function of ComplEx for given embeddings.

The embeddings have to be in a broadcastable shape.

Parameters
  • h (FloatTensor) – Head embeddings.

  • r (FloatTensor) – Relation embeddings.

  • t (FloatTensor) – Tail embeddings.

Return type

FloatTensor

Returns

shape: (…) The scores.

loss_default

alias of pykeen.losses.SoftplusLoss

loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'reduction': 'mean'}

The default parameters for the default loss function class

regularizer_default

alias of pykeen.regularizers.LpRegularizer

regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'normalize': True, 'p': 2.0, 'weight': 0.01}

The LP settings used by [trouillon2016] for ComplEx.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

class pykeen.models.ComplExLiteral(triples_factory, embedding_dim=50, automatic_memory_optimization=None, input_dropout=0.2, loss=None, preferred_device=None, random_seed=None)[source]

An implementation of ComplexLiteral from [agustinus2018] based on the LCWA training approach.

Initialize the model.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 300, 'low': 50, 'q': 50, 'type': <class 'int'>}, 'input_dropout': {'high': 0.3, 'low': 0.1, 'type': <class 'float'>}}

The default strategy for optimizing the model’s hyper-parameters

loss_default

alias of torch.nn.modules.loss.BCELoss

loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {}

The default parameters for the default loss function class

score_t(doubles)[source]

Forward pass using right side (tail) prediction for training with the LCWA.

Return type

Tensor

class pykeen.models.ConvE(triples_factory, input_channels=None, output_channels=32, embedding_height=None, embedding_width=None, kernel_height=3, kernel_width=3, input_dropout=0.2, output_dropout=0.3, feature_map_dropout=0.2, embedding_dim=200, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, regularizer=None, apply_batch_normalization=True)[source]

An implementation of ConvE from [dettmers2018].

The default setting uses batch normalization. Batch normalization normalizes the output of the activation functions, in order to ensure that the weights of the NN don’t become imbalanced and to speed up training. However, batch normalization is not the only way to achieve more robust and effective training [1]. Therefore, we added the flag ‘apply_batch_normalization’ to turn batch normalization on/off (it’s turned on as default).

[1]: Santurkar, Shibani, et al. “How does batch normalization help optimization?.” Advances in Neural Information Processing Systems. 2018.

Example usage:

>>> # Step 1: Get triples
>>> from pykeen.datasets import Nations
>>> dataset = Nations(create_inverse_triples=True)
>>> # Step 2: Configure the model
>>> from pykeen.models import ConvE
>>> model = ConvE(
...     embedding_dim       = 200,
...     input_channels      = 1,
...     output_channels     = 32,
...     embedding_height    = 10,
...     embedding_width     = 20,
...     kernel_height       = 3,
...     kernel_width        = 3,
...     input_dropout       = 0.2,
...     feature_map_dropout = 0.2,
...     output_dropout      = 0.3,
...     preferred_device    = 'gpu',
... )
>>> # Step 3: Configure the loop
>>> from torch.optim import Adam
>>> optimizer = Adam(params=model.get_grad_params())
>>> from pykeen.training import LCWATrainingLoop
>>> training_loop = LCWATrainingLoop(model=model, optimizer=optimizer)
>>> # Step 4: Train
>>> losses = training_loop.train(num_epochs=5, batch_size=256)
>>> # Step 5: Evaluate the model
>>> from pykeen.evaluation import RankBasedEvaluator
>>> evaluator = RankBasedEvaluator()
>>> metric_result = evaluator.evaluate(model=model, mapped_triples=dataset.testing.mapped_triples, batch_size=8192)

Initialize the model.

bn0: Optional[torch.nn.modules.batchnorm.BatchNorm2d]

If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,L)

bn1: Optional[torch.nn.modules.batchnorm.BatchNorm2d]

If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,H,W)

hpo_default: ClassVar[Mapping[str, Any]] = {'feature_map_dropout': {'high': 1.0, 'low': 0.0, 'type': <class 'float'>}, 'input_dropout': {'high': 1.0, 'low': 0.0, 'type': <class 'float'>}, 'output_channels': {'high': 64, 'low': 16, 'type': <class 'int'>}, 'output_dropout': {'high': 1.0, 'low': 0.0, 'type': <class 'float'>}}

The default strategy for optimizing the model’s hyper-parameters

loss_default

alias of pykeen.losses.BCEAfterSigmoidLoss

loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {}

The default parameters for the default loss function class

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.ConvKB(triples_factory, hidden_dropout_rate=0.0, embedding_dim=200, automatic_memory_optimization=None, loss=None, preferred_device=None, num_filters=400, random_seed=None, regularizer=None)[source]

An implementation of ConvKB from [nguyen2018].

See also

Initialize the model.

To be consistent with the paper, pass entity and relation embeddings pre-trained from TransE.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 300, 'low': 50, 'q': 50, 'type': <class 'int'>}, 'hidden_dropout_rate': {'high': 0.9, 'low': 0.1, 'type': <class 'float'>}, 'num_filters': {'high': 500, 'low': 300, 'q': 50, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

regularizer_default

alias of pykeen.regularizers.LpRegularizer

regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'apply_only_once': True, 'normalize': True, 'p': 2.0, 'weight': 0.0005}

The LP settings used by [nguyen2018] for ConvKB.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

class pykeen.models.DistMult(triples_factory, embedding_dim=50, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of DistMult from [yang2014].

This model simplifies RESCAL by restricting matrices representing relations as diagonal matrices.

Note:
  • For FB15k, Yang et al. report 2 negatives per each positive.

See also

Initialize the model.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

static interaction_function(h, r, t)[source]

Evaluate the interaction function for given embeddings.

The embeddings have to be in a broadcastable shape.

WARNING: Does not ensure forward constraints.

Parameters
  • h (FloatTensor) – shape: (…, e) Head embeddings.

  • r (FloatTensor) – shape: (…, e) Relation embeddings.

  • t (FloatTensor) – shape: (…, e) Tail embeddings.

Return type

FloatTensor

Returns

shape: (…) The scores.

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

regularizer_default

alias of pykeen.regularizers.LpRegularizer

regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'normalize': True, 'p': 2.0, 'weight': 0.1}

The LP settings used by [yang2014] for DistMult

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.DistMultLiteral(triples_factory, embedding_dim=50, automatic_memory_optimization=None, input_dropout=0.0, loss=None, preferred_device=None, random_seed=None)[source]

An implementation of DistMultLiteral from [agustinus2018].

Initialize the module.

compute_mr_loss(positive_scores, negative_scores)[source]

Compute the mean ranking loss for the positive and negative scores.

Return type

Tensor

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}, 'input_dropout': {'high': 1.0, 'low': 0, 'type': <class 'float'>}}

The default strategy for optimizing the model’s hyper-parameters

loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'margin': 0.0}

The default parameters for the default loss function class

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction for training with the LCWA.

Return type

Tensor

class pykeen.models.ERMLP(triples_factory, embedding_dim=50, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, hidden_dim=None, regularizer=None)[source]

An implementation of ERMLP from [dong2014].

This model uses a neural network-based approach.

Initialize the model.

hidden_dim

The multi-layer perceptron consisting of an input layer with 3 * self.embedding_dim neurons, a hidden layer with self.embedding_dim neurons and output layer with one neuron. The input is represented by the concatenation embeddings of the heads, relations and tail embeddings.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.ERMLPE(triples_factory, hidden_dim=300, input_dropout=0.2, hidden_dropout=0.3, embedding_dim=200, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An extension of ERMLP proposed by [sharifzadeh2019].

This model uses a neural network-based approach similar to ERMLP and with slight modifications. In ERMLP, the model is:

\[f(h, r, t) = \textbf{w}^{T} g(\textbf{W} [\textbf{h}; \textbf{r}; \textbf{t}])\]

whereas in ERMPLE the model is:

\[f(h, r, t) = \textbf{t}^{T} f(\textbf{W} (g(\textbf{W} [\textbf{h}; \textbf{r}]))\]

including dropouts and batch-norms between each two hidden layers. ConvE can be seen as a special case of ERMLPE that contains the unnecessary inductive bias of convolutional filters. The aim of this model is to show that lifting this bias from ConvE (which simply leaves us with a modified ERMLP model), not only reduces the number of parameters but also improves performance.

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}, 'hidden_dim': {'high': 450, 'low': 50, 'q': 25, 'type': <class 'int'>}, 'hidden_dropout': {'high': 0.8, 'low': 0.0, 'q': 0.1, 'type': <class 'float'>}, 'input_dropout': {'high': 0.8, 'low': 0.0, 'q': 0.1, 'type': <class 'float'>}}

The default strategy for optimizing the model’s hyper-parameters

loss_default

alias of pykeen.losses.BCEAfterSigmoidLoss

loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {}

The default parameters for the default loss function class

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.HolE(triples_factory, embedding_dim=200, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of HolE [nickel2016].

This model uses circular correlation to compose head and tail embeddings to afterwards compute the inner product with a relation embedding.

Initialize the model.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

static interaction_function(h, r, t)[source]

Evaluate the interaction function for given embeddings.

The embeddings have to be in a broadcastable shape.

Parameters
  • h (FloatTensor) – shape: (batch_size, num_entities, d) Head embeddings.

  • r (FloatTensor) – shape: (batch_size, num_entities, d) Relation embeddings.

  • t (FloatTensor) – shape: (batch_size, num_entities, d) Tail embeddings.

Return type

FloatTensor

Returns

shape: (batch_size, num_entities) The scores.

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.KG2E(triples_factory, embedding_dim=50, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, dist_similarity=None, c_min=0.05, c_max=5.0, regularizer=None)[source]

An implementation of KG2E from [he2015].

This model represents entities and relations as multi-dimensional Gaussian distributions.

Each relation is represented as

R ~ N(mu_r, Sigma_r)

Each entity is represented as

E ~ N(mu_e, Sigma_e)

For scoring, we compare E = (H - T) with R using a similarity function on distributions (KL div, Expected Likelihood).

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'c_max': {'high': 10.0, 'low': 1.0, 'type': <class 'float'>}, 'c_min': {'high': 0.1, 'low': 0.01, 'scale': 'log', 'type': <class 'float'>}, 'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.NTN(triples_factory, embedding_dim=100, automatic_memory_optimization=None, num_slices=4, loss=None, preferred_device=None, random_seed=None, non_linearity=None, regularizer=None)[source]

An implementation of NTN from [socher2013].

In NTN, a bilinear tensor layer relates the two entity vectors across multiple dimensions.

Scoring function:

u_R.T . f(h.T . W_R^[1:k] . t + V_r . [h; t] + b_R)

where h.T . W_R^[1:k] . t denotes the bilinear tensor product.

Initialize the model.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}, 'num_slices': {'high': 4, 'low': 2, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

score_h(rt_batch, slice_size=None)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch, slice_size=None)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.ProjE(triples_factory, embedding_dim=50, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, inner_non_linearity=None, regularizer=None)[source]

An implementation of ProjE from [shi2017].

See also

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

loss_default

alias of torch.nn.modules.loss.BCEWithLogitsLoss

loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'reduction': 'mean'}

The default parameters for the default loss function class

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.RESCAL(triples_factory, embedding_dim=50, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of RESCAL from [nickel2011].

This model represents relations as matrices and models interactions between latent features.

See also

Initialize the model.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

regularizer_default

alias of pykeen.regularizers.LpRegularizer

regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'normalize': True, 'p': 2.0, 'weight': 10}

The LP settings used by [nickel2011] for for RESCAL

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

Return type

FloatTensor

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.RGCN(triples_factory, embedding_dim=500, automatic_memory_optimization=None, loss=None, predict_with_sigmoid=False, preferred_device=None, random_seed=None, num_bases_or_blocks=5, num_layers=2, use_bias=True, use_batch_norm=False, activation_cls=None, activation_kwargs=None, base_model=None, sparse_messages_slcwa=True, edge_dropout=0.4, self_loop_dropout=0.2, edge_weighting=<function inverse_indegree_edge_weights>, decomposition='basis', buffer_messages=True)[source]

An implementation of R-GCN from [schlichtkrull2018].

This model uses graph convolutions with relation-specific weights.

Initialize the module.

activations: Optional[torch.nn.modules.container.ModuleList]

Activations for each layer (if used)

att: Optional[torch.nn.modules.container.ParameterList]

The relation-specific weights for each base shape: (num_relations, num_bases)

base_model: pykeen.models.base.EntityEmbeddingModel

Interaction model used as decoder

bases: Optional[torch.nn.modules.container.ParameterList]

The base weight matrices to generate relation-specific weights shape: (num_bases, embedding_dim, embedding_dim)

batch_norms: Optional[torch.nn.modules.container.ModuleList]

Batch normalization for each layer (if used)

biases: Optional[torch.nn.modules.container.ParameterList]

The biases for each layer (if used) shape of each element: (embedding_dim,)

blocks: Optional[torch.nn.modules.container.ParameterList]

The blocks of the relation-specific weight matrices shape: (num_relations, num_blocks, embedding_dim//num_blocks, embedding_dim//num_blocks)

hpo_default: ClassVar[Mapping[str, Any]] = {'activation_cls': {'choices': [None, <class 'torch.nn.modules.activation.ReLU'>, <class 'torch.nn.modules.activation.LeakyReLU'>], 'type': 'categorical'}, 'base_model_cls': {'choices': [<class 'pykeen.models.unimodal.distmult.DistMult'>, <class 'pykeen.models.unimodal.complex.ComplEx'>, <class 'pykeen.models.unimodal.ermlp.ERMLP'>], 'type': 'categorical'}, 'decomposition': {'choices': ['basis', 'block'], 'type': 'categorical'}, 'edge_dropout': {'high': 0.9, 'low': 0.0, 'type': <class 'float'>}, 'edge_weighting': {'choices': [None, <function inverse_indegree_edge_weights>, <function inverse_outdegree_edge_weights>, <function symmetric_edge_weights>], 'type': 'categorical'}, 'embedding_dim': {'high': 1000, 'low': 50, 'q': 50, 'type': <class 'int'>}, 'num_bases_or_blocks': {'high': 20, 'low': 2, 'q': 1, 'type': <class 'int'>}, 'num_layers': {'high': 5, 'low': 1, 'q': 1, 'type': <class 'int'>}, 'self_loop_dropout': {'high': 0.9, 'low': 0.0, 'type': <class 'float'>}, 'use_batch_norm': {'type': 'bool'}, 'use_bias': {'type': 'bool'}}

The default strategy for optimizing the model’s hyper-parameters

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

class pykeen.models.RotatE(triples_factory, embedding_dim=200, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of RotatE from [sun2019].

This model uses models relations as cotations in complex plane.

See also

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 1000, 'low': 125, 'q': 100, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

static interaction_function(h, r, t)[source]

Evaluate the interaction function of ComplEx for given embeddings.

The embeddings have to be in a broadcastable shape.

WARNING: No forward constraints are applied.

Parameters
  • h (FloatTensor) – shape: (…, e, 2) Head embeddings. Last dimension corresponds to (real, imag).

  • r (FloatTensor) – shape: (…, e, 2) Relation embeddings. Last dimension corresponds to (real, imag).

  • t (FloatTensor) – shape: (…, e, 2) Tail embeddings. Last dimension corresponds to (real, imag).

Return type

FloatTensor

Returns

shape: (…) The scores.

post_parameter_update()[source]

Normalize the length of relation vectors, if the forward constraint has not been applied yet.

Absolute value of complex number

\[|a + ib| = \sqrt{a^2 + b^2}\]

L2 norm of complex vector:

\[\|x\|^2 = \sum_{i=1}^d |x_i|^2 = \sum_{i=1}^d (x_i.re^2 + x_i.im^2) = (\sum_{i=1}^d x_i.re^2) + (\sum_{i=1}^d x_i.im^2) = \|x.re\|^2 + \|x.im\|^2 = \| [x.re; x.im] \|^2\]
Return type

None

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.SimplE(triples_factory, embedding_dim=200, automatic_memory_optimization=None, loss=None, preferred_device=None, random_seed=None, regularizer=None, clamp_score=None)[source]

An implementation of SimplE [kazemi2018].

This model extends CP by updating a triple, and the inverse triple.

See also

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

loss_default

alias of pykeen.losses.SoftplusLoss

loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {}

The default parameters for the default loss function class

regularizer_default

alias of pykeen.regularizers.PowerSumRegularizer

regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'normalize': True, 'p': 2.0, 'weight': 20}

The power sum settings used by [trouillon2016] for SimplE

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.StructuredEmbedding(triples_factory, embedding_dim=50, automatic_memory_optimization=None, scoring_fct_norm=1, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of Structured Embedding (SE) from [bordes2011].

This model projects different matrices for each relation head and tail entity.

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}, 'scoring_fct_norm': {'high': 2, 'low': 1, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

score_h(rt_batch, slice_size=None)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch, slice_size=None)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.TransD(triples_factory, embedding_dim=50, automatic_memory_optimization=None, relation_dim=30, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of TransD from [ji2015].

This model extends TransR to use fewer parameters.

See also

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 300, 'low': 20, 'q': 50, 'type': <class 'int'>}, 'relation_dim': {'high': 300, 'low': 20, 'q': 50, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

static interaction_function(h, h_p, r, r_p, t, t_p)[source]

Evaluate the interaction function for given embeddings.

The embeddings have to be in a broadcastable shape.

Parameters
  • h (FloatTensor) – shape: (batch_size, num_entities, d_e) Head embeddings.

  • h_p (FloatTensor) – shape: (batch_size, num_entities, d_e) Head projections.

  • r (FloatTensor) – shape: (batch_size, num_entities, d_r) Relation embeddings.

  • r_p (FloatTensor) – shape: (batch_size, num_entities, d_r) Relation projections.

  • t (FloatTensor) – shape: (batch_size, num_entities, d_e) Tail embeddings.

  • t_p (FloatTensor) – shape: (batch_size, num_entities, d_e) Tail projections.

Return type

FloatTensor

Returns

shape: (batch_size, num_entities) The scores.

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.TransE(triples_factory, embedding_dim=50, automatic_memory_optimization=None, scoring_fct_norm=1, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of TransE from [bordes2013].

This model considers a relation as a translation from the head to the tail entity.

See also

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 300, 'low': 50, 'q': 50, 'type': <class 'int'>}, 'scoring_fct_norm': {'high': 2, 'low': 1, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.TransH(triples_factory, embedding_dim=50, automatic_memory_optimization=None, scoring_fct_norm=1, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of TransH [wang2014].

This model extends TransE by applying the translation from head to tail entity in a relational-specific hyperplane.

See also

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 300, 'low': 50, 'q': 50, 'type': <class 'int'>}, 'scoring_fct_norm': {'high': 2, 'low': 1, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

regularize_if_necessary()[source]

Update the regularizer’s term given some tensors, if regularization is requested.

Return type

None

regularizer_default

alias of pykeen.regularizers.TransHRegularizer

regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'epsilon': 1e-05, 'weight': 0.05}

The settings used by [wang2014] for TransH

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.TransR(triples_factory, embedding_dim=50, automatic_memory_optimization=None, relation_dim=30, scoring_fct_norm=1, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of TransR from [lin2015].

This model extends TransE and TransH by considering different vector spaces for entities and relations.

Constraints:
  • $||h||_2 <= 1$: Done

  • $||r||_2 <= 1$: Done

  • $||t||_2 <= 1$: Done

  • $||h*M_r||_2 <= 1$: Done

  • $||t*M_r||_2 <= 1$: Done

Initialize the model.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 300, 'low': 20, 'q': 50, 'type': <class 'int'>}, 'relation_dim': {'high': 300, 'low': 20, 'q': 50, 'type': <class 'int'>}, 'scoring_fct_norm': {'high': 2, 'low': 1, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

static interaction_function(h, r, t, m_r)[source]

Evaluate the interaction function for given embeddings.

The embeddings have to be in a broadcastable shape.

Parameters
  • h (FloatTensor) – shape: (batch_size, num_entities, d_e) Head embeddings.

  • r (FloatTensor) – shape: (batch_size, num_entities, d_r) Relation embeddings.

  • t (FloatTensor) – shape: (batch_size, num_entities, d_e) Tail embeddings.

  • m_r (FloatTensor) – shape: (batch_size, num_entities, d_e, d_r) The relation specific linear transformations.

Return type

FloatTensor

Returns

shape: (batch_size, num_entities) The scores.

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.TuckER(triples_factory, embedding_dim=200, automatic_memory_optimization=None, relation_dim=None, loss=None, preferred_device=None, random_seed=None, dropout_0=0.3, dropout_1=0.4, dropout_2=0.5, regularizer=None, apply_batch_normalization=True)[source]

An implementation of TuckEr from [balazevic2019].

This model uses the Tucker tensor factorization.

See also

Initialize the model.

The dropout values correspond to the following dropouts in the model’s score function:

DO_2(BN(DO_0(BN(h)) x_1 DO_1(W x_2 r))) x_3 t

where h,r,t are the head, relation, and tail embedding, W is the core tensor, x_i denotes the tensor product along the i-th mode, BN denotes batch normalization, and DO dropout.

hpo_default: ClassVar[Mapping[str, Any]] = {'dropout_0': {'high': 0.4, 'low': 0.1, 'type': <class 'float'>}, 'dropout_1': {'high': 0.5, 'low': 0.1, 'type': <class 'float'>}, 'dropout_2': {'high': 0.6, 'low': 0.1, 'type': <class 'float'>}, 'embedding_dim': {'high': 300, 'low': 50, 'q': 50, 'type': <class 'int'>}, 'relation_dim': {'high': 200, 'low': 30, 'q': 25, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

loss_default

alias of pykeen.losses.BCEAfterSigmoidLoss

loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {}

The default parameters for the default loss function class

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

class pykeen.models.UnstructuredModel(triples_factory, embedding_dim=50, automatic_memory_optimization=None, scoring_fct_norm=1, loss=None, preferred_device=None, random_seed=None, regularizer=None)[source]

An implementation of Unstructured Model (UM) from [bordes2014].

Initialize the module.

hpo_default: ClassVar[Mapping[str, Any]] = {'embedding_dim': {'high': 350, 'low': 50, 'q': 25, 'type': <class 'int'>}, 'scoring_fct_norm': {'high': 2, 'low': 1, 'type': <class 'int'>}}

The default strategy for optimizing the model’s hyper-parameters

score_h(rt_batch)[source]

Forward pass using left side (head) prediction.

This method calculates the score for all possible heads for each (relation, tail) pair.

Parameters

rt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_hrt(hrt_batch)[source]

Forward pass.

This method takes head, relation and tail of each triple and calculates the corresponding score.

Parameters

hrt_batch (LongTensor) – torch.Tensor, shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, 1), dtype: float The score for each triple.

score_t(hr_batch)[source]

Forward pass using right side (tail) prediction.

This method calculates the score for all possible tails for each (head, relation) pair.

Parameters

hr_batch (LongTensor) – torch.Tensor, shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.

Return type

FloatTensor

Returns

torch.Tensor, shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

pykeen.models.get_model_cls(query)[source]

Get the model class.

Return type

Type[Model]

pykeen.models.models: Mapping[str, Type[pykeen.models.base.Model]] = {'complex': <class 'pykeen.models.unimodal.complex.ComplEx'>, 'complexliteral': <class 'pykeen.models.multimodal.complex_literal.ComplExLiteral'>, 'conve': <class 'pykeen.models.unimodal.conv_e.ConvE'>, 'convkb': <class 'pykeen.models.unimodal.conv_kb.ConvKB'>, 'distmult': <class 'pykeen.models.unimodal.distmult.DistMult'>, 'distmultliteral': <class 'pykeen.models.multimodal.distmult_literal.DistMultLiteral'>, 'ermlp': <class 'pykeen.models.unimodal.ermlp.ERMLP'>, 'ermlpe': <class 'pykeen.models.unimodal.ermlpe.ERMLPE'>, 'hole': <class 'pykeen.models.unimodal.hole.HolE'>, 'kg2e': <class 'pykeen.models.unimodal.kg2e.KG2E'>, 'ntn': <class 'pykeen.models.unimodal.ntn.NTN'>, 'proje': <class 'pykeen.models.unimodal.proj_e.ProjE'>, 'rescal': <class 'pykeen.models.unimodal.rescal.RESCAL'>, 'rgcn': <class 'pykeen.models.unimodal.rgcn.RGCN'>, 'rotate': <class 'pykeen.models.unimodal.rotate.RotatE'>, 'simple': <class 'pykeen.models.unimodal.simple.SimplE'>, 'structuredembedding': <class 'pykeen.models.unimodal.structured_embedding.StructuredEmbedding'>, 'transd': <class 'pykeen.models.unimodal.trans_d.TransD'>, 'transe': <class 'pykeen.models.unimodal.trans_e.TransE'>, 'transh': <class 'pykeen.models.unimodal.trans_h.TransH'>, 'transr': <class 'pykeen.models.unimodal.trans_r.TransR'>, 'tucker': <class 'pykeen.models.unimodal.tucker.TuckER'>, 'unstructuredmodel': <class 'pykeen.models.unimodal.unstructured_model.UnstructuredModel'>}

A mapping of models’ names to their implementations