nshtrainer 0.1.0__py3-none-any.whl

nshtrainer/callbacks/interval.py

@@ -0,0 +1,322 @@
+from collections.abc import Callable
+from typing import Literal
+
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+from typing_extensions import override
+
+Split = Literal["train", "val", "test", "predict"]
+
+
+def _check_step(step: int, interval: int, skip_first: bool = False):
+    if step % interval != 0:
+        return False
+    if skip_first and step == 0:
+        return False
+    return True
+
+
+class StepIntervalCallback(Callback):
+    def __init__(
+        self,
+        function: Callable[[Trainer, LightningModule], None],
+        *,
+        interval: int,
+        skip_first: bool = False,
+        splits: list[Split] = ["train", "val", "test", "predict"],
+    ):
+        super().__init__()
+
+        self.function = function
+        self.interval = interval
+        self.skip_first = skip_first
+        self.splits = set(splits)
+
+    @override
+    def on_train_batch_start(
+        self,
+        trainer,
+        pl_module,
+        batch,
+        batch_idx,
+        dataloader_idx=1,
+    ):
+        if (
+            not _check_step(
+                trainer.global_step,
+                self.interval,
+                skip_first=self.skip_first,
+            )
+            or "train" not in self.splits
+        ):
+            return
+        self.function(trainer, pl_module)
+
+    @override
+    def on_validation_batch_start(
+        self,
+        trainer,
+        pl_module,
+        batch,
+        batch_idx,
+        dataloader_idx=1,
+    ):
+        if (
+            not _check_step(
+                trainer.global_step,
+                self.interval,
+                skip_first=self.skip_first,
+            )
+            or "val" not in self.splits
+        ):
+            return
+        self.function(trainer, pl_module)
+
+    @override
+    def on_test_batch_start(
+        self,
+        trainer,
+        pl_module,
+        batch,
+        batch_idx,
+        dataloader_idx=1,
+    ):
+        if (
+            not _check_step(
+                trainer.global_step,
+                self.interval,
+                skip_first=self.skip_first,
+            )
+            or "test" not in self.splits
+        ):
+            return
+        self.function(trainer, pl_module)
+
+    @override
+    def on_predict_batch_start(
+        self,
+        trainer,
+        pl_module,
+        batch,
+        batch_idx,
+        dataloader_idx=1,
+    ):
+        if (
+            not _check_step(
+                trainer.global_step,
+                self.interval,
+                skip_first=self.skip_first,
+            )
+            or "predict" not in self.splits
+        ):
+            return
+        self.function(trainer, pl_module)
+
+
+class EpochIntervalCallback(Callback):
+    def __init__(
+        self,
+        function: Callable[[Trainer, LightningModule], None],
+        *,
+        interval: int,
+        skip_first: bool = False,
+        splits: list[Split] = ["train", "val", "test", "predict"],
+    ):
+        super().__init__()
+
+        self.function = function
+        self.interval = interval
+        self.skip_first = skip_first
+        self.splits = set(splits)
+
+    @override
+    def on_train_epoch_start(self, trainer, pl_module):
+        if (
+            not _check_step(
+                trainer.current_epoch,
+                self.interval,
+                skip_first=self.skip_first,
+            )
+            or "train" not in self.splits
+        ):
+            return
+        self.function(trainer, pl_module)
+
+    @override
+    def on_validation_epoch_start(self, trainer, pl_module):
+        if (
+            not _check_step(
+                trainer.current_epoch,
+                self.interval,
+                skip_first=self.skip_first,
+            )
+            or "val" not in self.splits
+        ):
+            return
+        self.function(trainer, pl_module)
+
+    @override
+    def on_test_epoch_start(self, trainer, pl_module):
+        if (
+            not _check_step(
+                trainer.current_epoch,
+                self.interval,
+                skip_first=self.skip_first,
+            )
+            or "test" not in self.splits
+        ):
+            return
+        self.function(trainer, pl_module)
+
+    @override
+    def on_predict_epoch_start(self, trainer, pl_module):
+        if (
+            not _check_step(
+                trainer.current_epoch,
+                self.interval,
+                skip_first=self.skip_first,
+            )
+            or "predict" not in self.splits
+        ):
+            return
+        self.function(trainer, pl_module)
+
+
+class IntervalCallback(Callback):
+    def __init__(
+        self,
+        function: Callable[[Trainer, LightningModule], None],
+        *,
+        step_interval: int | None = None,
+        epoch_interval: int | None = None,
+        skip_first: bool = False,
+        splits: list[Split] = ["train", "val", "test", "predict"],
+    ):
+        super().__init__()
+
+        self.callback = None
+
+        if step_interval is not None:
+            self.callback = StepIntervalCallback(
+                function,
+                interval=step_interval,
+                splits=splits,
+                skip_first=skip_first,
+            )
+        elif epoch_interval is not None:
+            self.callback = EpochIntervalCallback(
+                function,
+                interval=epoch_interval,
+                splits=splits,
+                skip_first=skip_first,
+            )
+        else:
+            raise ValueError("Either step_interval or epoch_interval must be specified")
+
+    @override
+    def on_train_batch_start(
+        self,
+        trainer,
+        pl_module,
+        batch,
+        batch_idx,
+        dataloader_idx=1,
+    ):
+        if not isinstance(self.callback, StepIntervalCallback):
+            return
+
+        self.callback.on_train_batch_start(
+            trainer,
+            pl_module,
+            batch,
+            batch_idx,
+            dataloader_idx=1,
+        )
+
+    @override
+    def on_validation_batch_start(
+        self,
+        trainer,
+        pl_module,
+        batch,
+        batch_idx,
+        dataloader_idx=1,
+    ):
+        if not isinstance(self.callback, StepIntervalCallback):
+            return
+
+        self.callback.on_validation_batch_start(
+            trainer,
+            pl_module,
+            batch,
+            batch_idx,
+            dataloader_idx=1,
+        )
+
+    @override
+    def on_test_batch_start(
+        self,
+        trainer,
+        pl_module,
+        batch,
+        batch_idx,
+        dataloader_idx=1,
+    ):
+        if not isinstance(self.callback, StepIntervalCallback):
+            return
+
+        self.callback.on_test_batch_start(
+            trainer,
+            pl_module,
+            batch,
+            batch_idx,
+            dataloader_idx=1,
+        )
+
+    @override
+    def on_predict_batch_start(
+        self,
+        trainer,
+        pl_module,
+        batch,
+        batch_idx,
+        dataloader_idx=1,
+    ):
+        if not isinstance(self.callback, StepIntervalCallback):
+            return
+
+        self.callback.on_predict_batch_start(
+            trainer,
+            pl_module,
+            batch,
+            batch_idx,
+            dataloader_idx=1,
+        )
+
+    @override
+    def on_train_epoch_start(self, trainer, pl_module):
+        if not isinstance(self.callback, EpochIntervalCallback):
+            return
+
+        self.callback.on_train_epoch_start(trainer, pl_module)
+
+    @override
+    def on_validation_epoch_start(self, trainer, pl_module):
+        if not isinstance(self.callback, EpochIntervalCallback):
+            return
+
+        self.callback.on_validation_epoch_start(trainer, pl_module)
+
+    @override
+    def on_test_epoch_start(self, trainer, pl_module):
+        if not isinstance(self.callback, EpochIntervalCallback):
+            return
+
+        self.callback.on_test_epoch_start(trainer, pl_module)
+
+    @override
+    def on_predict_epoch_start(self, trainer, pl_module):
+        if not isinstance(self.callback, EpochIntervalCallback):
+            return
+
+        self.callback.on_predict_epoch_start(trainer, pl_module)

nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -0,0 +1,45 @@
+import logging
+from pathlib import Path
+
+from lightning.fabric.utilities.types import _PATH
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Checkpoint
+from typing_extensions import override
+
+log = logging.getLogger(__name__)
+
+
+class LatestEpochCheckpoint(Checkpoint):
+    DEFAULT_FILENAME = "latest_epoch{epoch:02d}_step{step:04d}.ckpt"
+
+    def __init__(
+        self,
+        dirpath: _PATH,
+        filename: str | None = None,
+        save_weights_only: bool = False,
+    ):
+        super().__init__()
+
+        self._dirpath = Path(dirpath)
+        self._filename = filename or self.DEFAULT_FILENAME
+        self._save_weights_only = save_weights_only
+
+        # Also, we hold a reference to the last checkpoint path
+        # to be able to remove it when a new checkpoint is saved.
+        self._last_ckpt_path: Path | None = None
+
+    def _ckpt_path(self, trainer: Trainer):
+        return self._dirpath / self._filename.format(
+            epoch=trainer.current_epoch, step=trainer.global_step
+        )
+
+    @override
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        # Remove the last checkpoint if it exists
+        if self._last_ckpt_path is not None:
+            trainer.strategy.remove_checkpoint(self._last_ckpt_path)
+
+        # Save the new checkpoint
+        filepath = self._ckpt_path(trainer)
+        trainer.save_checkpoint(filepath, self._save_weights_only)
+        self._last_ckpt_path = filepath

nshtrainer/callbacks/log_epoch.py

@@ -0,0 +1,35 @@
+import logging
+import math
+from typing import Any
+
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+from typing_extensions import override
+
+log = logging.getLogger(__name__)
+
+
+class LogEpochCallback(Callback):
+    def __init__(self, metric_name: str = "computed_epoch"):
+        super().__init__()
+
+        self.metric_name = metric_name
+
+    @override
+    def on_train_batch_start(
+        self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int
+    ):
+        if trainer.logger is None:
+            return
+
+        # If trainer.num_training_batches is not set or is nan/inf, we cannot calculate the epoch
+        if (
+            not trainer.num_training_batches
+            or math.isnan(trainer.num_training_batches)
+            or math.isinf(trainer.num_training_batches)
+        ):
+            log.warning("Trainer has no valid num_training_batches. Cannot log epoch.")
+            return
+
+        epoch = pl_module.global_step / trainer.num_training_batches
+        pl_module.log(self.metric_name, epoch, on_step=True, on_epoch=False)

nshtrainer/callbacks/norm_logging.py

@@ -0,0 +1,187 @@
+from logging import getLogger
+from typing import Literal, cast
+
+import torch
+import torch.nn as nn
+from lightning.pytorch import Callback, LightningModule, Trainer
+from torch.optim import Optimizer
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+
+log = getLogger(__name__)
+
+
+def grad_norm(
+    module: nn.Module,
+    norm_type: float | int | str,
+    group_separator: str = "/",
+    grad: bool = True,
+) -> dict[str, torch.Tensor | float]:
+    """Compute each parameter's gradient's norm and their overall norm.
+
+    The overall norm is computed over all gradients together, as if they
+    were concatenated into a single vector.
+
+    Args:
+        module: :class:`torch.nn.Module` to inspect.
+        norm_type: The type of the used p-norm, cast to float if necessary.
+            Can be ``'inf'`` for infinity norm.
+        group_separator: The separator string used by the logger to group
+            the gradients norms in their own subfolder instead of the logs one.
+
+    Return:
+        norms: The dictionary of p-norms of each parameter's gradient and
+            a special entry for the total p-norm of the gradients viewed
+            as a single vector.
+    """
+    norm_type = float(norm_type)
+    if norm_type <= 0:
+        raise ValueError(
+            f"`norm_type` must be a positive number or 'inf' (infinity norm). Got {norm_type}"
+        )
+
+    if grad:
+        norms = {
+            f"grad_{norm_type}_norm{group_separator}{name}": p.grad.data.norm(norm_type)
+            for name, p in module.named_parameters()
+            if p.grad is not None
+        }
+        if norms:
+            total_norm = torch.tensor(list(norms.values())).norm(norm_type)
+            norms[f"grad_{norm_type}_norm_total"] = total_norm
+    else:
+        norms = {
+            f"param_{norm_type}_norm{group_separator}{name}": p.data.norm(norm_type)
+            for name, p in module.named_parameters()
+            if p.grad is not None
+        }
+        if norms:
+            total_norm = torch.tensor(list(norms.values())).norm(norm_type)
+            norms[f"param_{norm_type}_norm_total"] = total_norm
+
+    return norms
+
+
+def _to_norm_type(log_grad_norm_per_param: float | str | Literal[True]):
+    norm_type = 2.0
+    if log_grad_norm_per_param is not True:
+        norm_type = log_grad_norm_per_param
+    return norm_type
+
+
+def compute_norm(
+    pl_module: LightningModule,
+    optimizer: Optimizer | None = None,
+    p: float | str = 2.0,
+    *,
+    grad: bool,
+) -> torch.Tensor:
+    if optimizer is not None:
+        tensors = [
+            cast(torch.Tensor, p.grad if grad else p)
+            for group in optimizer.param_groups
+            for p in group["params"]
+            if p.grad is not None
+        ]
+    else:
+        tensors = [
+            p.grad if grad else p for p in pl_module.parameters() if p.grad is not None
+        ]
+
+    if not tensors:
+        return torch.tensor(0.0, device=pl_module.device)
+
+    return torch.norm(torch.stack([torch.norm(g, p=p) for g in tensors]), p=p)
+
+
+class NormLoggingCallback(Callback):
+    def __init__(self, config: "NormLoggingConfig"):
+        super().__init__()
+
+        self.config = config
+
+    def _perform_norm_logging(
+        self,
+        pl_module: LightningModule,
+        optimizer: Optimizer,
+        prefix: str,
+    ):
+        # Gradient norm logging
+        if log_grad_norm := self.config.log_grad_norm:
+            norm = compute_norm(
+                pl_module,
+                optimizer,
+                _to_norm_type(log_grad_norm),
+                grad=True,
+            )
+            pl_module.log(f"{prefix}grad_norm", norm, on_step=True, on_epoch=False)
+        if log_grad_norm_per_param := self.config.log_grad_norm_per_param:
+            norm_type = _to_norm_type(log_grad_norm_per_param)
+            pl_module.log_dict(
+                {
+                    f"{prefix}{k}": v
+                    for k, v in grad_norm(pl_module, norm_type, grad=True).items()
+                }
+            )
+
+        # Parameter norm logging
+        if log_param_norm := self.config.log_param_norm:
+            norm = compute_norm(
+                pl_module,
+                optimizer,
+                _to_norm_type(log_param_norm),
+                grad=False,
+            )
+            pl_module.log(f"{prefix}param_norm", norm, on_step=True, on_epoch=False)
+        if log_param_norm_per_param := self.config.log_param_norm_per_param:
+            norm_type = _to_norm_type(log_param_norm_per_param)
+            pl_module.log_dict(
+                {
+                    f"{prefix}{k}": v
+                    for k, v in grad_norm(pl_module, norm_type, grad=False).items()
+                }
+            )
+
+    @override
+    def on_after_backward(self, trainer: Trainer, pl_module: LightningModule):
+        if len(trainer.optimizers) == 1:
+            optimizer = trainer.optimizers[0]
+            self._perform_norm_logging(pl_module, optimizer, prefix="train/")
+        else:
+            for i, optimizer in enumerate(trainer.optimizers):
+                self._perform_norm_logging(
+                    pl_module, optimizer, prefix=f"train/optimizer_{i}/"
+                )
+
+
+class NormLoggingConfig(CallbackConfigBase):
+    name: Literal["norm_logging"] = "norm_logging"
+
+    log_grad_norm: bool | str | float = False
+    """If enabled, will log the gradient norm (averaged across all model parameters) to the logger."""
+    log_grad_norm_per_param: bool | str | float = False
+    """If enabled, will log the gradient norm for each model parameter to the logger."""
+
+    log_param_norm: bool | str | float = False
+    """If enabled, will log the parameter norm (averaged across all model parameters) to the logger."""
+    log_param_norm_per_param: bool | str | float = False
+    """If enabled, will log the parameter norm for each model parameter to the logger."""
+
+    def __bool__(self):
+        return any(
+            v
+            for v in (
+                self.log_grad_norm,
+                self.log_grad_norm_per_param,
+                self.log_param_norm,
+                self.log_param_norm_per_param,
+            )
+        )
+
+    @override
+    def construct_callbacks(self, root_config):
+        if not self:
+            return
+
+        yield NormLoggingCallback(self)

nshtrainer/callbacks/on_exception_checkpoint.py

@@ -0,0 +1,44 @@
+import datetime
+import logging
+import os
+from typing import Any
+
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import OnExceptionCheckpoint as _OnExceptionCheckpoint
+from typing_extensions import override
+
+log = logging.getLogger(__name__)
+
+
+class OnExceptionCheckpoint(_OnExceptionCheckpoint):
+    @property
+    @override
+    def ckpt_path(self) -> str:
+        ckpt_path = super().ckpt_path
+
+        # Remve the extension and add the current timestamp
+        timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        ckpt_path, ext = os.path.splitext(ckpt_path)
+        return f"{ckpt_path}_{timestamp}{ext}"
+
+    @override
+    def on_exception(self, trainer: Trainer, *_: Any, **__: Any) -> None:
+        # We override this to checkpoint the model manually,
+        # without calling the dist barrier.
+
+        # trainer.save_checkpoint(self.ckpt_path)
+
+        if trainer.model is None:
+            raise AttributeError(
+                "Saving a checkpoint is only possible if a model is attached to the Trainer. Did you call"
+                " `Trainer.save_checkpoint()` before calling `Trainer.{fit,validate,test,predict}`?"
+            )
+        checkpoint = trainer._checkpoint_connector.dump_checkpoint(weights_only=False)
+        trainer.strategy.save_checkpoint(
+            checkpoint, self.ckpt_path, storage_options=None
+        )
+        # self.strategy.barrier("Trainer.save_checkpoint") # <-- This is disabled
+
+    @override
+    def teardown(self, trainer: Trainer, *_: Any, **__: Any) -> None:
+        trainer.strategy.remove_checkpoint(self.ckpt_path)