From 38b79708003791682200ef15574c03a46d21f0d2 Mon Sep 17 00:00:00 2001 From: Kilian Lieret Date: Sat, 9 Mar 2024 20:23:33 -0500 Subject: [PATCH] Add LR monitor --- scripts/full_detector/continue_gc.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/scripts/full_detector/continue_gc.py b/scripts/full_detector/continue_gc.py index 7efe731..c8a8547 100644 --- 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)],