Refactor callbacks (#296)

* remove configure_callbacks in favor of specifying them in the config * remove callbacs from model * patience --> null
ML4GW · Oct 23, 2024 · b039e1b · b039e1b
1 parent ab6aa8a
commit b039e1b
Show file tree

Hide file tree

Showing 9 changed files with 977 additions and 1,238 deletions.
diff --git a/projects/data/poetry.lock b/projects/data/poetry.lock
diff --git a/projects/export/poetry.lock b/projects/export/poetry.lock
diff --git a/projects/infer/poetry.lock b/projects/infer/poetry.lock
diff --git a/projects/plots/poetry.lock b/projects/plots/poetry.lock
diff --git a/projects/train/config.yaml b/projects/train/config.yaml
@@ -28,8 +28,6 @@ model:
     weight_decay: 0.0
     learning_rate: 0.000585
     pct_lr_ramp: 0.115
-    # early stop
-    patience: null
 data:
   class_path: train.data.supervised.TimeDomainSupervisedAframeDataset
   init_args:
@@ -83,6 +81,23 @@ trainer:
       init_args:
         # save_dir:
         flush_logs_every_n_steps: 10
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.EarlyStopping
+      init_args:
+        monitor: "valid_auroc"
+        mode: "max"
+        patience: null
+    # custom model checkpoint for saving and
+    # tracing best model at end of traiing
+    # that will be used for downstream export
+    - class_path: train.callbacks.ModelCheckpoint
+      init_args:
+        monitor: "valid_auroc"
+        mode: "max"
+        save_top_k: 1
+        save_last: true
+        auto_insert_metric_name: false
+    - class_path: train.callbacks.SaveAugmentedBatch
   # uncomment below if you want to profile
   # profiler:
     # class_path: lightning.pytorch.profilers.PyTorchProfiler

diff --git a/projects/train/configs/bbh.yaml b/projects/train/configs/bbh.yaml
@@ -28,8 +28,6 @@ model:
     weight_decay: 0.0
     learning_rate: 0.000585
     pct_lr_ramp: 0.115
-    # early stop
-    patience: 50
 data:
   class_path: train.data.supervised.TimeDomainSupervisedAframeDataset
   init_args:
@@ -85,6 +83,23 @@ trainer:
         # save_dir:
         flush_logs_every_n_steps: 10
         name: "train_logs"
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.EarlyStopping
+      init_args:
+        monitor: "valid_auroc"
+        mode: "max"
+        patience: null
+    # custom model checkpoint for saving and
+    # tracing best model at end of traiing
+    # that will be used for downstream export
+    - class_path: train.callbacks.ModelCheckpoint
+      init_args:
+        monitor: "valid_auroc"
+        mode: "max"
+        save_top_k: 1
+        save_last: true
+        auto_insert_metric_name: false
+    - class_path: train.callbacks.SaveAugmentedBatch
   # uncomment below if you want to profile
   # profiler:
     # class_path: lightning.pytorch.profilers.PyTorchProfiler

diff --git a/projects/train/configs/bns.yaml b/projects/train/configs/bns.yaml
@@ -29,7 +29,7 @@ model:
     learning_rate: 0.000585
     pct_lr_ramp: 0.115
     # early stop
-    patience: 50
+    patience: null
 data:
   class_path: train.data.supervised.SpectrogramDomainSupervisedAframeDataset
   init_args:
@@ -87,6 +87,24 @@ trainer:
         # save_dir:
         flush_logs_every_n_steps: 10
         name: "train_logs"
+
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.EarlyStopping
+      init_args:
+        monitor: "valid_auroc"
+        mode: "max"
+        patience: 50
+    # custom model checkpoint for saving and
+    # tracing best model at end of traiing
+    # that will be used for downstream export
+    - class_path: train.callbacks.ModelCheckpoint
+      init_args:
+        monitor: "valid_auroc"
+        mode: "max"
+        save_top_k: 1
+        save_last: true
+        auto_insert_metric_name: false
+    - class_path: train.callbacks.SaveAugmentedBatch
   # uncomment below if you want to profile
   # profiler:
     # class_path: lightning.pytorch.profilers.PyTorchProfiler