Merge branch 'fix/sanity_check_metrics' into 'main'

Skip metric extraction when sanity checking See merge request es/ai/hannah/hannah!400
ekut-es · Aug 12, 2024 · e9530ee · e9530ee
2 parents 2ac38b6 + 90b281d
commit e9530ee
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/hannah/callbacks/optimization.py b/hannah/callbacks/optimization.py
@@ -112,6 +112,8 @@ def on_train_batch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModu
                 monitor_val = callback_metrics[monitor] * direction
                 if monitor.startswith("train"):
                     self._curves[monitor][trainer.global_step] = monitor_val
+
+                    self.values[monitor] = monitor_val
 
     def on_test_end(self, trainer, pl_module):
         """
@@ -151,6 +153,10 @@ def on_validation_end(self, trainer, pl_module):
         Returns:
 
         """
+        # Skip evaluation of validation metrics during sanity check
+        if trainer.sanity_checking:
+            return
+
         callback_metrics = trainer.callback_metrics
 
         for k, v in callback_metrics.items():
@@ -162,11 +168,8 @@ def on_validation_end(self, trainer, pl_module):
                 try:
                     monitor_val = float(callback_metrics[monitor])
                     directed_monitor_val = monitor_val * direction
-                    if (
-                        monitor not in self.values
-                        or directed_monitor_val < self.values[monitor]
-                    ):
-                        self.values[monitor] = directed_monitor_val
+
+                    self.values[monitor] = directed_monitor_val
                     self._curves[monitor][trainer.global_step] = directed_monitor_val
                 except Exception:
                     pass