🎲 🔮 Uncertainty Estimate via MC dropout (pykeen#688)

* add get_dropout_modules utility

* add uncertainty estimate via MC dropout

* move uncertain prediction to predict

* add predict_h_uncertain

* add tests

* add predict_{h,r}_uncertain

* fix method name

* add test for predict_{h,r}_uncertain

* run black

* trigger ci

* update __all__

trigger ci

* use named tuple

trigger ci

* add example use to docstring

trigger ci

* add missing info to docstring

trigger ci

* Small refactoring

* Add documentation todos

* Reorganize tests

* Clean up docs

* Reorganize

* Update uncertainty.py

* Update docs

* More reorg

* Update uncertainty.rst

* update some docstrings

* make helper public again

* rename file

* extend module docstring and fix std vs. var

* Update references.rst

* Update uncertainty.py

* add example to module docstring

* remove unused reference

* fix flake8

trigger ci

* Update uncertainty.py

* Delete uncertainty.rst

* Update index.rst

* trigger ci

* trigger ci

* Update uncertainty.py

* Trigger CI

Co-authored-by: Charles Tapley Hoyt <[email protected]>
Co-authored-by: PyKEEN_bot <[email protected]>

Author: 3 people

docs/source/index.rst

@@ -56,6 +56,7 @@ PyKEEN
    reference/ablation
    reference/lookup
    reference/predict
+   reference/uncertainty
    reference/sealant
    reference/constants
    reference/nn/index

docs/source/reference/uncertainty.rst

@@ -0,0 +1,6 @@
+Uncertainty
+===========
+.. automodapi:: pykeen.models.uncertainty
+    :no-heading:
+    :headings: --
+    :no-inheritance-diagram:

docs/source/references.rst

@@ -83,3 +83,6 @@ References
 
 .. [sharifzadeh2019vrd] Sharifzadeh, S., *et al*. (2019). `Improving Visual Relation Detection using Depth
    Maps <http://arxiv.org/abs/1905.00966>`_. *arXiv*, 1905.00966.
+
+.. [gal2016] Gal, Y., & Ghahramani, Z. (2016). `Dropout as a Bayesian Approximation: Representing Model Uncertainty in
+   Deep Learning <https://dl.acm.org/doi/10.5555/3045390.3045502>`_. ICML 2016.

src/pykeen/models/base.py

@@ -351,10 +351,10 @@ def predict_t(
 
         .. note::
 
-            We only expect the right side-side predictions, i.e., $(h,r,*)$ to change its
+            We only expect the right 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 :func:`predict_scores_all_heads()` has different behavior depending on
+            the :func:`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.
         """

src/pykeen/models/predict.py

@@ -258,7 +258,7 @@ def get_all_prediction_df(
 
     Example usage:
 
-    .. code-block:: python
+    .. code-block::
 
         from pykeen.pipeline import pipeline
         from pykeen.models.predict import get_all_prediction_df
@@ -667,16 +667,6 @@ def predict_triples_df(
     """
     Predict on labeled or mapped triples.
 
-    Example:
-    >>> from pykeen.pipeline import pipeline
-    >>> result = pipeline(dataset="nations", model="TransE")
-    >>> from pykeen.models.predict import predict_triples_df
-    >>> df = predict_triples_df(
-    ...     model=result.model,
-    ...     triples=("uk", "conferences", "brazil"),
-    ...     triples_factory=result.training,
-    ... )
-
     :param model:
         The model.
     :param triples: shape: (num_triples, 3)
@@ -699,6 +689,17 @@ def predict_triples_df(
 
     :raises ValueError:
         If label-based triples have been provided, but the triples factory does not provide a mapping.
+
+    The TransE model can be trained and used to predict a given triple.
+
+    >>> from pykeen.pipeline import pipeline
+    >>> result = pipeline(dataset="nations", model="TransE")
+    >>> from pykeen.models.predict import predict_triples_df
+    >>> df = predict_triples_df(
+    ...     model=result.model,
+    ...     triples=("uk", "conferences", "brazil"),
+    ...     triples_factory=result.training,
+    ... )
     """
     if triples is None:
         if triples_factory is None: