validate on best checkpoint after training

configs/train.yaml

@@ -50,6 +50,10 @@ train: True
 # lightning chooses best weights based on the metric specified in checkpoint callback
 test: False
 
+# evaluate on validation set, using best model weights achieved during training
+# lightning chooses best weights based on the metric specified in checkpoint callback
+validate: True
+
 # metric with mode (minimize or maximize) to monitor for checkpointing and early stopping callbacks
 monitor_metric: "val/f1"
 monitor_mode: "max"

src/datamodules/datamodule.py

@@ -79,6 +79,8 @@ def setup(self, stage: str):
 
         if stage == "fit":
             split_names = [self.train_split, self.val_split]
+        elif stage == "validate":
+            split_names = [self.val_split]
         elif stage == "test":
             split_names = [self.test_split]
         else:

src/train.py

@@ -190,6 +190,17 @@ def train(cfg: DictConfig) -> Tuple[dict, dict]:
         else:
             log.warning("the model is not saved because no save_dir is specified")
 
+    if cfg.get("validate"):
+        log.info("Starting validation!")
+        if best_ckpt_path == "":
+            log.warning("Best ckpt not found! Using current weights for validation...")
+        trainer.validate(model=model, datamodule=datamodule, ckpt_path=best_ckpt_path or None)
+    elif cfg.get("train"):
+        log.warning(
+            "Validation after training is skipped! That means, the finally reported validation scores are "
+            "the values from the *last* checkpoint, not from the *best* checkpoint (which is saved)!"
+        )
+
     if cfg.get("test"):
         log.info("Starting testing!")
         if best_ckpt_path == "":