predict_hrt_uncertain

predict_hrt_uncertain(model, hrt_batch, num_samples=5, *, mode=None)[source]

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

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

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

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

  • 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

UncertainPrediction

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.

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.predict 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],
)