Add LR monitor

gnn-tracking · Mar 10, 2024 · 38b7970 · 38b7970
1 parent 30218a8
commit 38b7970
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/scripts/full_detector/continue_gc.py b/scripts/full_detector/continue_gc.py
@@ -3,7 +3,12 @@
 from gnn_tracking.utils.loading import TrackingDataModule
 from gnn_tracking.utils.nomenclature import random_trial_name
 from lightning_fabric.plugins.environments.slurm import SLURMEnvironment
-from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, RichProgressBar
+from pytorch_lightning.callbacks import (
+    EarlyStopping,
+    LearningRateMonitor,
+    ModelCheckpoint,
+    RichProgressBar,
+)
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 from wandb_osh.lightning_hooks import TriggerWandbSyncLightningCallback
 
@@ -38,6 +43,7 @@ def cli_main():
                 ExpandWandbConfig(),
                 EarlyStopping(monitor="total", mode="min", patience=10),
                 ModelCheckpoint(save_top_k=2, monitor="total", mode="min"),
+                LearningRateMonitor(logging_interval="step", log_momentum=True),
             ],
             "logger": [tb_logger, wandb_logger],
             "plugins": [SLURMEnvironment(auto_requeue=False)],