predict_hrt_uncertain

predict_hrt_uncertain(model: Model, hrt_batch: Tensor, num_samples: int = 5, *, mode: Literal['training', 'validation', 'testing'] | None = None) → UncertainPrediction[source]

Calculate the scores with uncertainty quantification via Monte-Carlo dropout.

Parameters:

  • model (Model) – the model used for predicting scores

  • hrt_batch (Tensor) – shape: (number of triples, 3) The indices of (head, relation, tail) triples.

  • num_samples (int) – >1 the number of samples to draw

  • mode (Literal['training', 'validation', 'testing'] | None) – The pass mode, which is None in the transductive setting and one of “training”, “validation”, or “testing” in the inductive setting.

Returns:

shape: (number of triples, 1) The score for each triple, and an uncertainty score, where larger scores correspond to less certain predictions.

This function delegates to predict_uncertain_helper() by using pykeen.models.Model.score_hrt() as the score_method.

Return type:

UncertainPrediction

Warning

This function sets the model to evaluation mode and all dropout layers to training mode.

Example Usage:

from pykeen.pipeline import pipeline
from pykeen.models.uncertainty import predict_hrt_uncertain

result = pipeline(dataset="nations", model="ERMLPE")
prediction_with_uncertainty = predict_hrt_uncertain(
    model=result.model,
    hrt_batch=result.training.mapped_triples[0:8],
)