Model¶

class Model(triples_factory, loss=None, predict_with_sigmoid=False, preferred_device=None, random_seed=None)[source]¶

Bases: torch.nn.modules.module.Module, abc.ABC

A base module for KGE models.

Subclasses of Model can decide however they want on how to store entities’ and relations’ representations, how they want to be looked up, and how they should be scored. The OModel provides a commonly used interface for models storing entity and relation representations in the form of pykeen.nn.Embedding.

Initialize the module.

Parameters

triples_factory (TriplesFactory) – The triples factory facilitates access to the dataset.
loss (Optional[Loss]) – The loss to use. If None is given, use the loss default specific to the model subclass.
predict_with_sigmoid (bool) – Whether to apply sigmoid onto the scores when predicting scores. Applying sigmoid at prediction time may lead to exactly equal scores for certain triples with very high, or very low score. When not trained with applying sigmoid (or using BCEWithLogitsLoss), the scores are not calibrated to perform well with sigmoid.
preferred_device (Union[None, str, device]) – The preferred device for model training and inference.
random_seed (Optional[int]) – A random seed to use for initialising the model’s weights. Should be set when aiming at reproducibility.
regularizer – A regularizer to use for training.

Attributes Summary

`can_slice_h`	Whether score_h supports slicing.
`can_slice_r`	Whether score_r supports slicing.
`can_slice_t`	Whether score_t supports slicing.
`loss_default_kwargs`	The default parameters for the default loss function class
`num_entities`	The number of entities in the knowledge graph.
`num_parameter_bytes`	Calculate the number of bytes used for all parameters of the model.
`num_relations`	The number of unique relation types in the knowledge graph.

Methods Summary

`compute_loss`(tensor_1, tensor_2)	Compute the loss for functions requiring two separate tensors as input.
`get_all_prediction_df`(*[, k, batch_size])	Compute scores for all triples, optionally returning only the k highest scoring.
`get_grad_params`()	Get the parameters that require gradients.
`get_head_prediction_df`(relation_label, …)	Predict heads for the given relation and tail (given by label).
`get_relation_prediction_df`(head_label, …)	Predict relations for the given head and tail (given by label).
`get_tail_prediction_df`(head_label, …)	Predict tails for the given head and relation (given by label).
`load_state`(path)	Load the state of the model.
`post_forward_pass`()	Run after calculating the forward loss.
`post_parameter_update`()	Has to be called after each parameter update.
`predict_h`(rt_batch[, slice_size])	Forward pass using left side (head) prediction for obtaining scores of all possible heads.
`predict_hrt`(hrt_batch)	Calculate the scores for triples.
`predict_r`(ht_batch[, slice_size])	Forward pass using middle (relation) prediction for obtaining scores of all possible relations.
`predict_t`(hr_batch[, slice_size])	Forward pass using right side (tail) prediction for obtaining scores of all possible tails.
`reset_parameters_`()	Reset all parameters of the model and enforce model constraints.
`save_state`(path)	Save the state of the model.
`score_h`(rt_batch)	Forward pass using left side (head) prediction.
`score_h_inverse`(rt_batch[, slice_size])	Score all heads for a batch of (r,t)-pairs using the tail predictions for the inverses \((t,r_{inv},*)\).
`score_hrt`(hrt_batch)	Forward pass.
`score_hrt_inverse`(hrt_batch)	Score triples based on inverse triples, i.e., compute \(f(h,r,t)\) based on \(f(t,r_{inv},h)\).
`score_r`(ht_batch)	Forward pass using middle (relation) prediction.
`score_t`(hr_batch)	Forward pass using right side (tail) prediction.
`score_t_inverse`(hr_batch[, slice_size])	Score all tails for a batch of (h,r)-pairs using the head predictions for the inverses \((*,r_{inv},h)\).
`to_device_`()	Transfer model to device.

Attributes Documentation

can_slice_h¶

Whether score_h supports slicing.

Return type: bool

can_slice_r¶

Whether score_r supports slicing.

Return type: bool

can_slice_t¶

Whether score_t supports slicing.

Return type: bool

loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = {'margin': 1.0, 'reduction': 'mean'}¶: The default parameters for the default loss function class

num_entities¶

The number of entities in the knowledge graph.

Return type: int

num_parameter_bytes¶

Calculate the number of bytes used for all parameters of the model.

Return type: int

num_relations¶

The number of unique relation types in the knowledge graph.

Return type: int

Methods Documentation

abstract compute_loss(tensor_1, tensor_2)[source]¶

Compute the loss for functions requiring two separate tensors as input.

Parameters

tensor_1 (FloatTensor) – shape: s The tensor containing predictions or positive scores.
tensor_2 (FloatTensor) – shape: s The tensor containing target values or the negative scores.

Return type

FloatTensor

Returns

dtype: float, scalar The label loss value.

Note

generally the two tensors do not need to have the same shape, but only one which is broadcastable.

get_all_prediction_df(*, k=None, batch_size=1, **kwargs)[source]¶

Compute scores for all triples, optionally returning only the k highest scoring.

Note

This operation is computationally very expensive for reasonably-sized knowledge graphs.

Warning

Setting k=None may lead to huge memory requirements.

Parameters

k (Optional[int]) – The number of triples to return. Set to None, to keep all.
batch_size (int) – The batch size to use for calculating scores.
kwargs – Additional kwargs to pass to pykeen.models.predict.get_all_prediction_df().

Return type

Union[ScorePack, DataFrame]

Returns

shape: (k, 3) A tensor containing the k highest scoring triples, or all possible triples if k=None.

get_grad_params()[source]¶

Get the parameters that require gradients.

Return type: Iterable[Parameter]

get_head_prediction_df(relation_label, tail_label, **kwargs)[source]¶

Predict heads for the given relation and tail (given by label).

Parameters

relation_label (str) – The string label for the relation
tail_label (str) – The string label for the tail entity
kwargs – Keyword arguments passed to pykeen.models.predict.get_head_prediction_df()

The following example shows that after you train a model on the Nations dataset, you can score all entities w.r.t a given relation and tail entity.

>>> from pykeen.pipeline import pipeline
>>> result = pipeline(
...     dataset='Nations',
...     model='RotatE',
... )
>>> df = result.model.get_head_prediction_df('accusation', 'brazil')

Return type: DataFrame

get_relation_prediction_df(head_label, tail_label, **kwargs)[source]¶

Predict relations for the given head and tail (given by label).

Parameters

head_label (str) – The string label for the head entity
tail_label (str) – The string label for the tail entity
kwargs – Keyword arguments passed to pykeen.models.predict.get_relation_prediction_df()

Return type

DataFrame

get_tail_prediction_df(head_label, relation_label, **kwargs)[source]¶

Predict tails for the given head and relation (given by label).

Parameters

head_label (str) – The string label for the head entity
relation_label (str) – The string label for the relation
kwargs – Keyword arguments passed to pykeen.models.predict.get_tail_prediction_df()

The following example shows that after you train a model on the Nations dataset, you can score all entities w.r.t a given head entity and relation.

>>> from pykeen.pipeline import pipeline
>>> result = pipeline(
...     dataset='Nations',
...     model='RotatE',
... )
>>> df = result.model.get_tail_prediction_df('brazil', 'accusation')

Return type: DataFrame

load_state(path)[source]¶

Load the state of the model.

Parameters: path (str) – Path of the file where to load the state from.
Return type: None

post_forward_pass()[source]¶: Run after calculating the forward loss.

post_parameter_update()[source]¶

Has to be called after each parameter update.

Return type: None

predict_h(rt_batch, slice_size=None)[source]¶

Forward pass using left side (head) prediction for obtaining scores of all possible heads.

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

Note

If the model has been trained with inverse relations, the task of predicting the head entities becomes the task of predicting the tail entities of the inverse triples, i.e., \(f(*,r,t)\) is predicted by means of \(f(t,r_{inv},*)\).

Additionally, the model is set to evaluation mode.

Parameters

rt_batch (LongTensor) – shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.
slice_size (Optional[int]) – >0 The divisor for the scoring function when using slicing.

Return type

FloatTensor

Returns

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

predict_hrt(hrt_batch)[source]¶

Calculate the scores for triples.

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

Additionally, the model is set to evaluation mode.

Parameters: hrt_batch (LongTensor) – shape: (number of triples, 3), dtype: long The indices of (head, relation, tail) triples.
Return type: FloatTensor
Returns: shape: (number of triples, 1), dtype: float The score for each triple.

predict_r(ht_batch, slice_size=None)[source]¶

Forward pass using middle (relation) prediction for obtaining scores of all possible relations.

This method calculates the score for all possible relations for each (head, tail) pair.

Additionally, the model is set to evaluation mode.

Parameters

ht_batch (LongTensor) – shape: (batch_size, 2), dtype: long The indices of (head, tail) pairs.
slice_size (Optional[int]) – >0 The divisor for the scoring function when using slicing.

Return type

FloatTensor

Returns

shape: (batch_size, num_relations), dtype: float For each h-t pair, the scores for all possible relations.

predict_t(hr_batch, slice_size=None)[source]¶

Forward pass using right side (tail) prediction for obtaining scores of all possible tails.

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

Additionally, the model is set to evaluation mode.

Parameters

hr_batch (LongTensor) – shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.
slice_size (Optional[int]) – >0 The divisor for the scoring function when using slicing.

Return type

FloatTensor

Returns

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

Note

We only expect the right side-side predictions, i.e., \((h,r,*)\) to change its default behavior when the model has been trained with inverse relations (mainly because of the behavior of the LCWA training approach). This is why the predict_scores_all_heads() has different behavior depending on if inverse triples were used in training, and why this function has the same behavior regardless of the use of inverse triples.

reset_parameters_()[source]¶: Reset all parameters of the model and enforce model constraints.

save_state(path)[source]¶

Save the state of the model.

Parameters: path (str) – Path of the file where to store the state in.
Return type: None

abstract 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) – shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs.
Return type: FloatTensor
Returns: shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads.

score_h_inverse(rt_batch, slice_size=None)[source]¶: Score all heads for a batch of (r,t)-pairs using the tail predictions for the inverses \((t,r_{inv},*)\).

abstract 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) – shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples.
Raises: NotImplementedError – If the method was not implemented for this class.
Return type: FloatTensor
Returns: shape: (batch_size, 1), dtype: float The score for each triple.

score_hrt_inverse(hrt_batch)[source]¶

Score triples based on inverse triples, i.e., compute \(f(h,r,t)\) based on \(f(t,r_{inv},h)\).

When training with inverse relations, the model produces two (different) scores for a triple \((h,r,t) \in K\). The forward score is calculated from \(f(h,r,t)\) and the inverse score is calculated from \(f(t,r_{inv},h)\). This function enables users to inspect the scores obtained by using the corresponding inverse triples.

Return type: FloatTensor

abstract score_r(ht_batch)[source]¶

Forward pass using middle (relation) prediction.

This method calculates the score for all possible relations for each (head, tail) pair.

Parameters: ht_batch (LongTensor) – shape: (batch_size, 2), dtype: long The indices of (head, tail) pairs.
Return type: FloatTensor
Returns: shape: (batch_size, num_relations), dtype: float For each h-t pair, the scores for all possible relations.

abstract 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) – shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs.
Return type: FloatTensor
Returns: shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails.

score_t_inverse(hr_batch, slice_size=None)[source]¶: Score all tails for a batch of (h,r)-pairs using the head predictions for the inverses \((*,r_{inv},h)\).

to_device_()[source]¶: Transfer model to device.