improve: log early stopping best score

Author: samedii

lantern/early_stopping.py

@@ -1,3 +1,4 @@
+import torch.utils.tensorboard
 from typing import Optional
 
 from lantern import FunctionalBase
@@ -6,9 +7,14 @@
 class EarlyStopping(FunctionalBase):
     """Keeps track of the best score and how long ago it was calculated."""
 
+    tensorboard_logger: torch.utils.tensorboard.SummaryWriter
     best_score: Optional[float] = None
     scores_since_improvement: int = -1
 
+    class Config:
+        arbitrary_types_allowed = True
+        allow_mutation = False
+
     def score(self, value):
         if self.best_score is None or value >= self.best_score:
             return self.replace(
@@ -31,5 +37,10 @@ def print(self):
         )
         return self
 
-    def log(self):
-        pass
+    def log(self, step):
+        self.tensorboard_logger.add_scalar(
+            "best_score",
+            self.best_score,
+            step,
+        )
+        return self

test/test_mnist.py

@@ -52,7 +52,7 @@ def test_mnist():
     )
 
     tensorboard_logger = torch.utils.tensorboard.SummaryWriter()
-    early_stopping = lantern.EarlyStopping()
+    early_stopping = lantern.EarlyStopping(tensorboard_logger=tensorboard_logger)
     gradient_metrics = lantern.Metrics(
         name="gradient",
         tensorboard_logger=tensorboard_logger,
@@ -108,7 +108,7 @@ def test_mnist():
             torch.save(optimizer.state_dict(), "optimizer.pt")
         elif early_stopping.scores_since_improvement > 5:
             break
-        early_stopping.print()
+        early_stopping.log(epoch).print()
 
 
 if __name__ == "__main__":