|
|
"""Rich progress bar callback with epoch-level train/val losses.""" |
|
|
|
|
|
from __future__ import annotations |
|
|
|
|
|
from typing import Any |
|
|
|
|
|
import torch |
|
|
from lightning.pytorch.callbacks import RichProgressBar |
|
|
|
|
|
|
|
|
class LossLoggingCallback(RichProgressBar): |
|
|
"""Show epoch-level mean train/val losses directly in Rich progress bar.""" |
|
|
|
|
|
def __init__(self) -> None: |
|
|
"""Initialize callback state used to aggregate losses per epoch.""" |
|
|
super().__init__() |
|
|
self._train_loss_sum = 0.0 |
|
|
self._train_loss_count = 0 |
|
|
self._val_loss_sum = 0.0 |
|
|
self._val_loss_count = 0 |
|
|
self._latest_train_loss: float | None = None |
|
|
self._latest_val_loss: float | None = None |
|
|
|
|
|
@staticmethod |
|
|
def _as_float(value: Any) -> float | None: |
|
|
"""Convert tensor/scalar values to float for logging.""" |
|
|
if value is None: |
|
|
return None |
|
|
if torch.is_tensor(value): |
|
|
return float(value.detach().cpu().item()) |
|
|
if isinstance(value, (int, float)): |
|
|
return float(value) |
|
|
return None |
|
|
|
|
|
@staticmethod |
|
|
def _extract_loss_from_outputs(outputs: Any) -> float | None: |
|
|
"""Extract loss scalar from Lightning step outputs.""" |
|
|
if outputs is None: |
|
|
return None |
|
|
if isinstance(outputs, dict): |
|
|
for key in ("loss", "train_loss", "val_loss"): |
|
|
loss = LossLoggingCallback._as_float(outputs.get(key)) |
|
|
if loss is not None: |
|
|
return loss |
|
|
loss = LossLoggingCallback._as_float(outputs) |
|
|
if loss is not None: |
|
|
return loss |
|
|
if hasattr(outputs, "loss"): |
|
|
return LossLoggingCallback._as_float(outputs.loss) |
|
|
return None |
|
|
|
|
|
@staticmethod |
|
|
def _extract_train_loss_fallback(metrics: dict[str, Any]) -> float | None: |
|
|
"""Fallback: get train loss from callback metrics.""" |
|
|
for key in ("train_loss_epoch", "train_loss", "train/loss"): |
|
|
loss = LossLoggingCallback._as_float(metrics.get(key)) |
|
|
if loss is not None: |
|
|
return loss |
|
|
return None |
|
|
|
|
|
@staticmethod |
|
|
def _extract_val_loss_fallback(metrics: dict[str, Any]) -> float | None: |
|
|
"""Fallback: get validation loss from callback metrics.""" |
|
|
for key in ("val_loss_epoch", "val_loss", "val/loss"): |
|
|
loss = LossLoggingCallback._as_float(metrics.get(key)) |
|
|
if loss is not None: |
|
|
return loss |
|
|
return None |
|
|
|
|
|
def on_train_epoch_start(self, trainer, pl_module) -> None: |
|
|
"""Reset train accumulators at epoch start.""" |
|
|
super().on_train_epoch_start(trainer, pl_module) |
|
|
self._train_loss_sum = 0.0 |
|
|
self._train_loss_count = 0 |
|
|
|
|
|
def on_validation_epoch_start(self, trainer, pl_module) -> None: |
|
|
"""Reset validation accumulators at epoch start.""" |
|
|
super().on_validation_epoch_start(trainer, pl_module) |
|
|
self._val_loss_sum = 0.0 |
|
|
self._val_loss_count = 0 |
|
|
|
|
|
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx) -> None: |
|
|
"""Accumulate per-batch train loss.""" |
|
|
super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx) |
|
|
loss = self._extract_loss_from_outputs(outputs) |
|
|
if loss is None: |
|
|
return |
|
|
self._train_loss_sum += loss |
|
|
self._train_loss_count += 1 |
|
|
|
|
|
def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0) -> None: |
|
|
"""Accumulate per-batch validation loss.""" |
|
|
super().on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) |
|
|
if trainer.sanity_checking: |
|
|
return |
|
|
loss = self._extract_loss_from_outputs(outputs) |
|
|
if loss is None: |
|
|
return |
|
|
self._val_loss_sum += loss |
|
|
self._val_loss_count += 1 |
|
|
|
|
|
def on_train_epoch_end(self, trainer, pl_module) -> None: |
|
|
"""Store mean train loss at train epoch end.""" |
|
|
super().on_train_epoch_end(trainer, pl_module) |
|
|
train_loss = ( |
|
|
self._train_loss_sum / self._train_loss_count |
|
|
if self._train_loss_count > 0 |
|
|
else None |
|
|
) |
|
|
if train_loss is None: |
|
|
metrics = dict(trainer.callback_metrics) |
|
|
train_loss = self._extract_train_loss_fallback(metrics) |
|
|
self._latest_train_loss = train_loss |
|
|
|
|
|
def on_validation_epoch_end(self, trainer, pl_module) -> None: |
|
|
"""Store mean val loss at validation epoch end.""" |
|
|
super().on_validation_epoch_end(trainer, pl_module) |
|
|
if trainer.sanity_checking: |
|
|
return |
|
|
|
|
|
val_loss = ( |
|
|
self._val_loss_sum / self._val_loss_count |
|
|
if self._val_loss_count > 0 |
|
|
else None |
|
|
) |
|
|
metrics = dict(trainer.callback_metrics) |
|
|
if val_loss is None: |
|
|
val_loss = self._extract_val_loss_fallback(metrics) |
|
|
self._latest_val_loss = val_loss |
|
|
|
|
|
def get_metrics(self, trainer, pl_module) -> dict[str, float]: |
|
|
"""Append epoch-level losses to Rich progress bar metrics.""" |
|
|
metrics = super().get_metrics(trainer, pl_module) |
|
|
if self._latest_train_loss is not None: |
|
|
metrics["train_loss"] = round(float(self._latest_train_loss), 6) |
|
|
if self._latest_val_loss is not None: |
|
|
metrics["val_loss"] = round(float(self._latest_val_loss), 6) |
|
|
return metrics |
|
|
|
