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 scoreshrt_batch (
LongTensor
) – shape: (number of triples, 3) The indices of (head, relation, tail) triples.num_samples (
int
) – >1 the number of samples to drawmode (
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:
- 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 usingpykeen.models.Model.score_hrt()
as thescore_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.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], )