Model

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

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

A base module for all of the KGE models.

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 (Optional[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

modules_not_supporting_sub_batching

Return all modules not supporting sub-batching.

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.

regularizer_default_kwargs

The default parameters for the default regularizer class

Methods Summary

compute_label_loss(predictions, labels)

Compute the classification loss.

compute_mr_loss(positive_scores, negative_scores)

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

compute_self_adversarial_negative_sampling_loss(…)

Compute self adversarial negative sampling loss.

get_grad_params()

Get the parameters that require gradients.

load_state(path)

Load the state of the model.

make_labeled_df(tensor, **kwargs)

Take a tensor of triples and make a pandas dataframe with labels.

post_parameter_update()

Has to be called after each parameter update.

predict_heads(relation_label, tail_label[, …])

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

predict_scores(triples)

Calculate the scores for triples.

predict_scores_all_heads(rt_batch[, slice_size])

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

predict_scores_all_relations(ht_batch[, …])

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

predict_scores_all_tails(hr_batch[, slice_size])

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

predict_tails(head_label, relation_label[, …])

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

regularize_if_necessary(*tensors)

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

reset_parameters_()

Reset all parameters of the model and enforce model constraints.

save_state(path)

Save the state of the model.

score_all_triples([k, batch_size, …])

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

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_cpu_()

Transfer the entire model to CPU.

to_device_()

Transfer model to device.

to_gpu_()

Transfer the entire model to GPU.

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

modules_not_supporting_sub_batching

Return all modules not supporting sub-batching.

Return type

Collection[Module]

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

regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None

The default parameters for the default regularizer class

Methods Documentation

compute_label_loss(predictions, labels)[source]

Compute the classification loss.

Parameters
  • predictions (FloatTensor) – shape: s The tensor containing predictions.

  • labels (FloatTensor) – shape: s The tensor containing labels.

Return type

FloatTensor

Returns

dtype: float, scalar The label loss value.

compute_mr_loss(positive_scores, negative_scores)[source]

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

Parameters
  • positive_scores (FloatTensor) – shape: s, dtype: float The scores for positive triples.

  • negative_scores (FloatTensor) – shape: s, dtype: float The scores for negative triples.

Raises

RuntimeError – If the chosen loss function does not allow the calculation of margin ranking

Return type

FloatTensor

Returns

dtype: float, scalar The margin ranking loss value.

compute_self_adversarial_negative_sampling_loss(positive_scores, negative_scores)[source]

Compute self adversarial negative sampling loss.

Parameters
  • positive_scores (FloatTensor) – shape: s The tensor containing the positive scores.

  • negative_scores (FloatTensor) – shape: s Tensor containing the negative scores.

Raises

RuntimeError – If the chosen loss does not allow the calculation of self adversarial negative sampling losses.

Return type

FloatTensor

Returns

dtype: float, scalar The loss value.

get_grad_params()[source]

Get the parameters that require gradients.

Return type

Iterable[Parameter]

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

make_labeled_df(tensor, **kwargs)[source]

Take a tensor of triples and make a pandas dataframe with labels.

Parameters
  • tensor (LongTensor) – shape: (n, 3) The triples, ID-based and in format (head_id, relation_id, tail_id).

  • kwargs (Union[Tensor, ndarray, Sequence]) – Any additional number of columns. Each column needs to be of shape (n,). Reserved column names: {“head_id”, “head_label”, “relation_id”, “relation_label”, “tail_id”, “tail_label”}.

Return type

DataFrame

Returns

A dataframe with n rows, and 6 + len(kwargs) columns.

post_parameter_update()[source]

Has to be called after each parameter update.

Return type

None

predict_heads(relation_label, tail_label, add_novelties=True, remove_known=False, testing=None)[source]

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

Parameters
  • relation_label (str) – The string label for the relation

  • tail_label (str) – The string label for the tail entity

  • add_novelties (bool) – Should the dataframe include a column denoting if the ranked head entities correspond to novel triples?

  • remove_known (bool) – Should non-novel triples (those appearing in the training set) be shown with the results? On one hand, this allows you to better assess the goodness of the predictions - you want to see that the non-novel triples generally have higher scores. On the other hand, if you’re doing hypothesis generation, they may pose as a distraction. If this is set to True, then non-novel triples will be removed and the column denoting novelty will be excluded, since all remaining triples will be novel. Defaults to false.

  • testing (Optional[LongTensor]) – The mapped_triples from the testing triples factory (TriplesFactory.mapped_triples)

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.predict_heads('accusation', 'brazil')
Return type

DataFrame

predict_scores(triples)[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

triples (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_scores_all_heads(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_scores_all_relations(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_scores_all_tails(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.

predict_tails(head_label, relation_label, add_novelties=True, remove_known=False, testing=None)[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

  • add_novelties (bool) – Should the dataframe include a column denoting if the ranked tail entities correspond to novel triples?

  • remove_known (bool) – Should non-novel triples (those appearing in the training set) be shown with the results? On one hand, this allows you to better assess the goodness of the predictions - you want to see that the non-novel triples generally have higher scores. On the other hand, if you’re doing hypothesis generation, they may pose as a distraction. If this is set to True, then non-novel triples will be removed and the column denoting novelty will be excluded, since all remaining triples will be novel. Defaults to false.

  • testing (Optional[LongTensor]) – The mapped_triples from the testing triples factory (TriplesFactory.mapped_triples)

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.predict_tails('brazil', 'accusation')
Return type

DataFrame

regularize_if_necessary(*tensors)[source]

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

Parameters

tensors (FloatTensor) – The tensors that should be passed to the regularizer to update its term.

Return type

None

reset_parameters_()[source]

Reset all parameters of the model and enforce model constraints.

Return type

Model

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

score_all_triples(k=None, batch_size=1, return_tensors=False, add_novelties=True, remove_known=False, testing=None)[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.

Return type

Union[Tuple[LongTensor, FloatTensor], DataFrame]

Returns

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

Example usage:

from pykeen.pipeline import pipeline

# Train a model (quickly)
result = pipeline(model='RotatE', dataset='Nations', training_kwargs=dict(num_epochs=5))
model = result.model

# Get scores for *all* triples
tensor = model.score_all_triples()
df = model.make_labeled_df(tensor)

# Get scores for top 15 triples
top_df = model.score_all_triples(k=15)
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

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.

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_cpu_()[source]

Transfer the entire model to CPU.

Return type

Model

to_device_()[source]

Transfer model to device.

Return type

Model

to_gpu_()[source]

Transfer the entire model to GPU.

Return type

Model