enhance overfit_batches logic + revert logging changes

courtois-neuromod · Jan 15, 2025 · e1c972b · e1c972b
1 parent f1f3be0
commit e1c972b
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/cneuromax/fitting/deeplearning/litmodule/base.py b/cneuromax/fitting/deeplearning/litmodule/base.py
@@ -135,8 +135,7 @@ def stage_step(
         if isinstance(data, list):
             data = tuple(data)
         loss: Num[Tensor, " *_"] = self.step(data, stage)
-        self.log(name=f"{stage}_step/loss", value=loss, on_step=True)
-        self.log(name=f"{stage}_epoch/loss", value=loss, on_epoch=True)
+        self.log(name=f"{stage}/loss", value=loss)
         return loss
 
     @final

diff --git a/cneuromax/fitting/deeplearning/train.py b/cneuromax/fitting/deeplearning/train.py
@@ -69,13 +69,14 @@ def train(
         >= TORCH_COMPILE_MINIMUM_CUDA_VERSION
     ):
         litmodule.nnmodule = torch.compile(litmodule.nnmodule)  # type: ignore [assignment]
-    litmodule.trainer = trainer
+    if trainer.overfit_batches > 0:
+        datamodule.val_dataloader = datamodule.train_dataloader
     trainer.fit(
         model=litmodule,
         datamodule=datamodule,
         ckpt_path=config.ckpt_path,
     )
     """TODO: Add logic for HPO"""
     return trainer.validate(model=litmodule, datamodule=datamodule)[0][
-        "val_epoch/loss"
+        "val/loss"
     ]
diff --git a/cneuromax/fitting/deeplearning/utils/lightning.py b/cneuromax/fitting/deeplearning/utils/lightning.py
@@ -65,7 +65,7 @@ def instantiate_trainer(
     callbacks.append(
         ModelCheckpoint(
             dirpath=trainer_partial.keywords["default_root_dir"],
-            monitor="val_epoch/loss",
+            monitor="val/loss",
             save_last=True,
             save_top_k=1,
             every_n_train_steps=save_every_n_train_steps,