Model
 class Model(*, triples_factory, loss=None, loss_kwargs=None, predict_with_sigmoid=False, random_seed=None)[source]

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. TheOModel
provides a commonly used interface for models storing entity and relation representations in the form ofpykeen.nn.Embedding
.Initialize the module.
 Parameters:
triples_factory (
KGInfo
) – The triples factory facilitates access to the dataset.loss (
Union
[str
,Loss
,Type
[Loss
],None
]) – The loss to use. If None is given, use the loss default specific to the model subclass.loss_kwargs (
Optional
[Mapping
[str
,Any
]]) – keywordbased parameters passed to the loss instance upon instantiationpredict_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.random_seed (
Optional
[int
]) – A random seed to use for initialising the model’s weights. Should be set when aiming at reproducibility.
Attributes Summary
Return the model's device.
The default parameters for the default loss function class
Calculate the number of bytes used for all parameters of the model.
Calculate the number of parameters of the model.
Return the real number of relations (without inverses).
Methods Summary
Get the regularization term for the loss function.
Get the parameters that require gradients.
load_state
(path)Load the state of the model.
Run after calculating the forward loss.
Has to be called after each parameter update.
predict
(hrt_batch, target[, full_batch, ids])Predict scores for the given target.
predict_h
(rt_batch, **kwargs)Forward pass using left side (head) prediction for obtaining scores of all possible heads.
predict_hrt
(hrt_batch, *[, mode])Calculate the scores for triples.
predict_r
(ht_batch, **kwargs)Forward pass using middle (relation) prediction for obtaining scores of all possible relations.
predict_t
(hr_batch, **kwargs)Forward pass using right side (tail) prediction for obtaining scores of all possible tails.
Reset all parameters of the model and enforce model constraints.
save_state
(path)Save the state of the model.
score_h
(rt_batch, *[, slice_size, mode, heads])Forward pass using left side (head) prediction.
score_h_inverse
(rt_batch, *[, heads])Score all heads for a batch of (r,t)pairs using the tail predictions for the inverses \((t,r_{inv},*)\).
score_hrt
(hrt_batch, *[, mode])Forward pass.
score_hrt_inverse
(hrt_batch, *[, mode])Score triples based on inverse triples, i.e., compute \(f(h,r,t)\) based on \(f(t,r_{inv},h)\).
score_r
(ht_batch, *[, slice_size, mode, ...])Forward pass using middle (relation) prediction.
score_t
(hr_batch, *[, slice_size, mode, tails])Forward pass using right side (tail) prediction.
score_t_inverse
(hr_batch, *[, tails])Score all tails for a batch of (h,r)pairs using the head predictions for the inverses \((*,r_{inv},h)\).
Attributes Documentation
 loss_default_kwargs: ClassVar[Mapping[str, Any]  None] = {'margin': 1.0, 'reduction': 'mean'}
The default parameters for the default loss function class
 num_parameter_bytes
Calculate the number of bytes used for all parameters of the model.
 Return type:
Methods Documentation
 abstract collect_regularization_term()[source]
Get the regularization term for the loss function.
 Return type:
FloatTensor
 predict(hrt_batch, target, full_batch=True, ids=None, **kwargs)[source]
Predict scores for the given target.
 Parameters:
hrt_batch (
LongTensor
) – shape: (batch_size, 3) or (batch_size, 2) the full batch, or the relevant part of ittarget (
Literal
[‘head’, ‘relation’, ‘tail’]) – the target to predictfull_batch (
bool
) – whether hrt_batch is the full batch, or only the “input” part of the target prediction methodids (
Optional
[LongTensor
]) – restrict prediction to only those idskwargs – additional keywordbased parameters passed to the specific target prediction method.
 Raises:
ValueError – if the target is invalid
 Return type:
FloatTensor
 Returns:
shape: (batch_size, num) the scores
 predict_h(rt_batch, **kwargs)[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.kwargs – additional keywordbased parameters passed to
Model.score_h()
 Return type:
FloatTensor
 Returns:
shape: (batch_size, num_heads), dtype: float For each rt pair, the scores for all possible heads.
 predict_hrt(hrt_batch, *, mode=None)[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:
 Return type:
FloatTensor
 Returns:
shape: (number of triples, 1), dtype: float The score for each triple.
 predict_r(ht_batch, **kwargs)[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.kwargs – additional keywordbased parameters passed to
Model.score_r()
 Return type:
FloatTensor
 Returns:
shape: (batch_size, num_relations), dtype: float For each ht pair, the scores for all possible relations.
 predict_t(hr_batch, **kwargs)[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.kwargs – additional keywordbased parameters passed to
Model.score_t()
 Return type:
FloatTensor
 Returns:
shape: (batch_size, num_tails), dtype: float For each hr pair, the scores for all possible tails.
Note
We only expect the right sidepredictions, 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_h()
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.
 abstract score_h(rt_batch, *, slice_size=None, mode=None, heads=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
) – 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.mode (
Optional
[Literal
[‘training’, ‘validation’, ‘testing’]]) – The pass mode, which is None in the transductive setting and one of “training”, “validation”, or “testing” in the inductive setting.heads (
Optional
[LongTensor
]) – shape: (num_heads,)  (batch_size, num_heads) head entity indices to score against. If None, scores against all entities (from the given mode).
 Return type:
FloatTensor
 Returns:
shape: (batch_size, num_heads), dtype: float For each rt pair, the scores for all possible heads.
 score_h_inverse(rt_batch, *, heads=None, **kwargs)[source]
Score all heads for a batch of (r,t)pairs using the tail predictions for the inverses \((t,r_{inv},*)\).
 Parameters:
rt_batch (LongTensor) –
heads (LongTensor  None) –
 abstract score_hrt(hrt_batch, *, mode=None)[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.mode (
Optional
[Literal
[‘training’, ‘validation’, ‘testing’]]) – The pass mode, which is None in the transductive setting and one of “training”, “validation”, or “testing” in the inductive setting.
 Return type:
FloatTensor
 Returns:
shape: (batch_size, 1), dtype: float The score for each triple.
 score_hrt_inverse(hrt_batch, *, mode=None)[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.
 abstract score_r(ht_batch, *, slice_size=None, mode=None, relations=None)[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.slice_size (
Optional
[int
]) – >0 The divisor for the scoring function when using slicing.mode (
Optional
[Literal
[‘training’, ‘validation’, ‘testing’]]) – The pass mode, which is None in the transductive setting and one of “training”, “validation”, or “testing” in the inductive setting.relations (
Optional
[LongTensor
]) – shape: (num_relations,)  (batch_size, num_relations) relation indices to score against. If None, scores against all relations (from the given mode).
 Return type:
FloatTensor
 Returns:
shape: (batch_size, num_real_relations), dtype: float For each ht pair, the scores for all possible relations.
 abstract score_t(hr_batch, *, slice_size=None, mode=None, tails=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
) – 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.mode (
Optional
[Literal
[‘training’, ‘validation’, ‘testing’]]) – The pass mode, which is None in the transductive setting and one of “training”, “validation”, or “testing” in the inductive setting.tails (
Optional
[LongTensor
]) – shape: (num_tails,)  (batch_size, num_tails) tail entity indices to score against. If None, scores against all entities (from the given mode).
 Return type:
FloatTensor
 Returns:
shape: (batch_size, num_tails), dtype: float For each hr pair, the scores for all possible tails.