EPIG implementation

baal/active/__init__.py

@@ -30,5 +30,6 @@ def get_heuristic(
         "variance": heuristics.Variance,
         "precomputed": heuristics.Precomputed,
         "batch_bald": heuristics.BatchBALD,
+        "epig": heuristics.EPIG,
     }[name](shuffle_prop=shuffle_prop, reduction=reduction, **kwargs)
     return heuristic

baal/active/active_loop.py

@@ -81,10 +81,21 @@ def step(self, pool=None) -> bool:
         if len(pool) > 0:
             if isinstance(self.heuristic, heuristics.Random):
                 probs = np.random.uniform(low=0, high=1, size=(len(pool), 1))
+                targets = None
             else:
                 probs = self.get_probabilities(pool, **self.kwargs)
+                if isinstance(self.heuristic, heuristics.EPIG):
+                    targets = self.get_probabilities(self.dataset, **self.kwargs)
+                else:
+                    targets = None
             if probs is not None and (isinstance(probs, types.GeneratorType) or len(probs) > 0):
-                to_label, uncertainty = self.heuristic.get_ranks(probs)
+                to_label, uncertainty = self.heuristic.get_ranks(probs, targets)
+                log.info(
+                    "Uncertainty",
+                    mean=uncertainty.mean(),
+                    std=uncertainty.std(),
+                    median=np.median(uncertainty),
+                )
                 if indices is not None:
                     to_label = indices[np.array(to_label)]
                 if self.uncertainty_folder is not None: