enable choosing from final or best ckpt

VMarsocci · Nov 8, 2024 · dc90879 · dc90879
1 parent ecfefa1
commit dc90879
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/configs/test.yaml b/configs/test.yaml
@@ -9,6 +9,7 @@ batch_size: 8
 
 
 # EXPERIMENT
+use_final_ckpt: false
 finetune: false
 ckpt_dir: ???
 

diff --git a/configs/train.yaml b/configs/train.yaml
@@ -19,6 +19,7 @@ limited_label_val: 1
 limited_label_strategy: stratified  # Options: stratified, oversampled, random
 stratification_bins: 3 #  number of bins for stratified sampling, only for stratified
 data_replicate: 1
+use_final_ckpt: false
 
 
 defaults:

diff --git a/pangaea/run.py b/pangaea/run.py
@@ -24,6 +24,7 @@
 from pangaea.utils.utils import (
     fix_seed,
     get_best_model_ckpt_path,
+    get_final_model_ckpt_path,
     get_generator,
     seed_worker,
 )
@@ -278,8 +279,12 @@ def main(cfg: DictConfig) -> None:
     test_evaluator: Evaluator = instantiate(
         cfg.task.evaluator, val_loader=test_loader, exp_dir=exp_dir, device=device
     )
-    best_model_ckpt_path = get_best_model_ckpt_path(exp_dir)
-    test_evaluator.evaluate(decoder, "best_model", best_model_ckpt_path)
+
+    if cfg.use_final_ckpt:
+        model_ckpt_path = get_final_model_ckpt_path(exp_dir)
+    else:
+        model_ckpt_path = get_best_model_ckpt_path(exp_dir)
+    test_evaluator.evaluate(decoder, "test_model", model_ckpt_path)
 
     if cfg.use_wandb and rank == 0:
         wandb.finish()