Track learning rate with callback. n_epochs redundant with the `epo…

…chs` tracked by lightning.

graphium/config/_loader.py

@@ -13,7 +13,7 @@
 
 # Lightning
 from lightning import Trainer
-from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
+from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint, LearningRateMonitor
 from lightning.pytorch.loggers import Logger, WandbLogger
 from loguru import logger
 
@@ -415,6 +415,11 @@ def load_trainer(
     if "model_checkpoint" in cfg_trainer.keys():
         callbacks.append(ModelCheckpoint(**cfg_trainer["model_checkpoint"]))
 
+    if "learning_rate_monitor" in cfg_trainer.keys():
+        callbacks.append(LearningRateMonitor(**cfg_trainer["learning_rate_monitor"]))
+    else:
+        callbacks.append(LearningRateMonitor())
+
     # Define the logger parameters
     wandb_cfg = config["constants"].get("wandb")
     if wandb_cfg is not None:

graphium/trainer/predictor.py

@@ -618,11 +618,6 @@ def on_validation_epoch_end(self) -> None:
         concatenated_metrics_logs = self.task_epoch_summary.concatenate_metrics_logs(metrics_logs)
         concatenated_metrics_logs["val/mean_time"] = torch.tensor(self.mean_val_time_tracker.mean_value)
         concatenated_metrics_logs["val/mean_tput"] = self.mean_val_tput_tracker.mean_value
-
-        if hasattr(self.optimizers(), "param_groups"):
-            lr = self.optimizers().param_groups[0]["lr"]
-            concatenated_metrics_logs["lr"] = torch.tensor(lr)
-        concatenated_metrics_logs["n_epochs"] = torch.tensor(self.current_epoch, dtype=torch.float32)
         self.log_dict(concatenated_metrics_logs)
 
         # Save yaml file with the per-task metrics summaries