Add Stopping Criteria for loop

baal/active/stopping_criteria.py

@@ -0,0 +1,63 @@
+from typing import List, Dict
+
+import numpy as np
+
+from baal import ActiveLearningDataset
+
+
+class StoppingCriterion:
+    def __init__(self, active_dataset: ActiveLearningDataset):
+        self._active_ds = active_dataset
+
+    def should_stop(self, metrics: Dict[str, float], uncertainty: List[float]) -> bool:
+        raise NotImplementedError
+
+
+class LabellingBudgetStoppingCriterion(StoppingCriterion):
+    """Stops when the labelling budget is exhausted."""
+
+    def __init__(self, active_dataset: ActiveLearningDataset, labelling_budget: int):
+        super().__init__(active_dataset)
+        self._start_length = len(active_dataset)
+        self.labelling_budget = labelling_budget
+
+    def should_stop(self, uncertainty: List[float]) -> bool:
+        return (len(self._active_ds) - self._start_length) >= self.labelling_budget
+
+
+class LowAverageUncertaintyStoppingCriterion(StoppingCriterion):
+    """Stops when the average uncertainty is on average below a threshold."""
+
+    def __init__(self, active_dataset: ActiveLearningDataset, avg_uncertainty_thresh: float):
+        super().__init__(active_dataset)
+        self.avg_uncertainty_thresh = avg_uncertainty_thresh
+
+    def should_stop(self, metrics: Dict[str, float], uncertainty: List[float]) -> bool:
+        return np.mean(uncertainty) < self.avg_uncertainty_thresh
+
+
+class EarlyStoppingCriterion(StoppingCriterion):
+    """Early stopping on a particular metrics.
+
+    Notes:
+    We don't have any mandatory dependency with an early stopping implementation.
+    So we have our own.
+    """
+
+    def __init__(
+        self,
+        active_dataset: ActiveLearningDataset,
+        metric_name: str,
+        patience: int = 10,
+        epsilon: float = 1e-4,
+    ):
+        super().__init__(active_dataset)
+        self.metric_name = metric_name
+        self.patience = patience
+        self.epsilon = epsilon
+        self._acc = []
+
+    def should_stop(self, metrics: Dict[str, float], uncertainty: List[float]) -> bool:
+        self._acc.append(metrics[self.metric_name])
+        near_threshold = np.isclose(np.array(self._acc), self._acc[-1], atol=self.epsilon)
+        return len(near_threshold) > self.patience and near_threshold[-self.patience].all()

experiments/mlp_mcdropout.py

@@ -9,6 +9,7 @@
 from baal import ActiveLearningDataset, ModelWrapper
 from baal.active import ActiveLearningLoop
 from baal.active.heuristics import BALD
+from baal.active.stopping_criteria import LabellingBudgetStoppingCriterion
 from baal.bayesian.dropout import patch_module
 
 use_cuda = torch.cuda.is_available()
@@ -54,16 +55,18 @@
 
 # Following Gal 2016, we reset the weights at the beginning of each step.
 initial_weights = deepcopy(model.state_dict())
+stopping_criterion = LabellingBudgetStoppingCriterion(
+    active_dataset=al_dataset, labelling_budget=10
+)
 
-for step in range(100):
+while True:
     model.load_state_dict(initial_weights)
     train_loss = wrapper.train_on_dataset(
         al_dataset, optimizer=optimizer, batch_size=32, epoch=10, use_cuda=use_cuda
     )
     test_loss = wrapper.test_on_dataset(test_ds, batch_size=32, use_cuda=use_cuda)
 
     pprint(wrapper.get_metrics())
-    flag = al_loop.step()
-    if not flag:
-        # We are done labelling! stopping
+    al_loop.step()
+    if stopping_criterion.should_stop():
         break