Repository: ElementAI/baal
Pull request #4 (2019-10-08 16:35:27)

Add Early Stopping to fit_and_test
## Summary:
Add Early Stopping to fit_and_test
### Features:
* Able to specify the patience before quitting.

## Checklist:

* [x] Your code is documented (To validate this, add your module to `tests/documentation_test.py`).
* [x] Your code is tested with unit tests.
* [ ] You moved your Issue to the PR state.



Patch:
diff --git a/src/baal/modelwrapper.py b/src/baal/modelwrapper.py
index 0395389..ad2128f 100644
--- a/src/baal/modelwrapper.py
+++ b/src/baal/modelwrapper.py
@@ -148,7 +148,9 @@ class ModelWrapper:
                                    use_cuda: bool,
                                    workers: int = 4,
                                    collate_fn: Optional[Callable] = None,
-                                   return_best_weights=False):
+                                   return_best_weights=False,
+                                   patience=None,
+                                   min_epoch_for_es=0):
         """
         Train and test the model on both Dataset `train_dataset`, `test_dataset`.
 
@@ -162,21 +164,31 @@ class ModelWrapper:
             workers (int): Number of workers to use.
             collate_fn (Optional[Callable]): The collate function to use.
             return_best_weights (bool): If True, will keep the best weights and return them.
+            patience (Optional[int]): If provided, will use early stopping to stop after
+                                        `patience` epoch without improvement.
+            min_epoch_for_es (int): Epoch at which the early stopping starts.
 
         Returns:
             History and best weights if required.
         """
         best_weight = None
         best_loss = 1e10
+        best_epoch = 0
         hist = []
-        for epoch in range(epoch):
+        for e in range(epoch):
             _ = self.train_on_dataset(train_dataset, optimizer, batch_size, 1,
                                       use_cuda, workers, collate_fn)
             te_loss = self.test_on_dataset(test_dataset, batch_size, use_cuda, workers, collate_fn)
             hist.append({k: v.value for k, v in self.metrics.items()})
-            if return_best_weights and te_loss < best_loss:
+            if te_loss < best_loss:
+                best_epoch = e
                 best_loss = te_loss
-                best_weight = deepcopy(self.state_dict())
+                if return_best_weights:
+                    best_weight = deepcopy(self.state_dict())
+
+            if patience is not None and (e - best_epoch) > patience and (e > min_epoch_for_es):
+                # Early stopping
+                break
 
         if return_best_weights:
             return hist, best_weight


Test patch:
diff --git a/tests/modelwrapper_test.py b/tests/modelwrapper_test.py
index 67dbb91..702d3c2 100644
--- a/tests/modelwrapper_test.py
+++ b/tests/modelwrapper_test.py
@@ -1,4 +1,5 @@
 import unittest
+from unittest.mock import Mock
 
 import numpy as np
 import pytest
@@ -278,6 +279,7 @@ class ModelWrapperTest(unittest.TestCase):
         assert (self.wrapper.metrics['test_cls_report'].value['accuracy'] != 0).any()
 
     def test_train_and_test(self):
+
         res = self.wrapper.train_and_test_on_datasets(self.dataset, self.dataset, self.optim, 32, 5,
                                                       False, return_best_weights=False)
         assert len(res) == 5
@@ -286,6 +288,26 @@ class ModelWrapperTest(unittest.TestCase):
         assert len(res) == 2
         assert len(res[0]) == 5
         assert isinstance(res[1], dict)
+        mock = Mock()
+        mock.side_effect = (((np.linspace(0, 50) - 10) / 10) ** 2).tolist()
+        self.wrapper.test_on_dataset = mock
+        res = self.wrapper.train_and_test_on_datasets(self.dataset, self.dataset,
+                                                      self.optim, 32, 50,
+                                                      False, return_best_weights=True, patience=1)
+
+        assert len(res) == 2
+        assert len(res[0]) < 50
+
+        mock = Mock()
+        mock.side_effect = (((np.linspace(0, 50) - 10) / 10) ** 2).tolist()
+        self.wrapper.test_on_dataset = mock
+        res = self.wrapper.train_and_test_on_datasets(self.dataset, self.dataset,
+                                                      self.optim, 32, 50,
+                                                      False, return_best_weights=True, patience=1,
+                                                      min_epoch_for_es=20)
+        assert len(res) == 2
+        assert len(res[0]) < 50 and len(res[0]) > 20
+
 
 
 if __name__ == '__main__':