|
from typing import Optional, Union
|
|
|
|
import lightning.pytorch as pl
|
|
import torch
|
|
from lightning import LightningModule, Trainer
|
|
from lightning.pytorch.callbacks import Callback
|
|
from torch import Tensor, nn
|
|
from torch.utils._foreach_utils import (
|
|
_group_tensors_by_device_and_dtype,
|
|
_has_foreach_support,
|
|
)
|
|
|
|
|
|
@torch.no_grad()
|
|
def grad_norm(
|
|
parameters: Union[Tensor, list[Tensor]],
|
|
norm_type: float = 2.0,
|
|
) -> float:
|
|
"""
|
|
Returns the norm of the gradients of the given parameters.
|
|
|
|
Args:
|
|
parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
|
|
single Tensor that will have gradients normalized
|
|
norm_type (float): type of the used p-norm.
|
|
|
|
Returns:
|
|
Total norm of the parameter gradients (viewed as a single vector).
|
|
"""
|
|
|
|
if isinstance(parameters, Tensor):
|
|
parameters = [parameters]
|
|
|
|
grads = [p.grad for p in parameters if p.grad is not None]
|
|
if len(grads) == 0:
|
|
return None
|
|
|
|
first_device = grads[0].device
|
|
grouped_grads: dict[
|
|
tuple[torch.device, torch.dtype], list[list[Tensor]]
|
|
] = _group_tensors_by_device_and_dtype(
|
|
[[g.detach() for g in grads]]
|
|
)
|
|
|
|
norms = []
|
|
for (device, _), ([grads], _) in grouped_grads.items():
|
|
if _has_foreach_support(grads, device=device):
|
|
norms.extend(torch._foreach_norm(grads, norm_type))
|
|
else:
|
|
norms.extend([torch.norm(g, norm_type) for g in grads])
|
|
|
|
return torch.norm(torch.stack([norm.to(first_device) for norm in norms]), norm_type)
|
|
|
|
|
|
class GradNormMonitor(Callback):
|
|
"""
|
|
Callback that computes the gradient norm of the model parameters.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
norm_type: float = 2.0,
|
|
logging_interval: str = "step",
|
|
sub_module: Optional[Union[str, list[str]]] = None,
|
|
) -> None:
|
|
"""
|
|
Args:
|
|
norm_type (float): type of the used p-norm.
|
|
logging_interval (str): "step" or "epoch".
|
|
"""
|
|
super().__init__()
|
|
|
|
self.norm_type = norm_type
|
|
self.logging_interval = logging_interval
|
|
self.sub_module = sub_module
|
|
|
|
def on_after_backward(self, trainer: Trainer, model: LightningModule) -> None:
|
|
"""
|
|
Computes the gradient norm of the model parameters and logs it to the logger.
|
|
|
|
Args:
|
|
trainer (Trainer): The trainer object
|
|
model (LightningModule): The current lightningModule
|
|
"""
|
|
|
|
lightning_model = model
|
|
|
|
if self.sub_module is None:
|
|
return self.log_sub_module_grad_norm(lightning_model, model, "")
|
|
|
|
sub_modules = self.sub_module
|
|
if isinstance(sub_modules, str):
|
|
sub_modules = [sub_modules]
|
|
|
|
for sub_module in sub_modules:
|
|
self.log_sub_module_grad_norm(
|
|
lightning_model, getattr(model, sub_module), f"/{sub_module}"
|
|
)
|
|
|
|
def log_sub_module_grad_norm(
|
|
self, lightning_model: LightningModule, model: nn.Module, path: str
|
|
) -> None:
|
|
grad_norm_val = grad_norm(model.parameters(), self.norm_type)
|
|
if grad_norm_val is None:
|
|
return
|
|
|
|
on_step = self.logging_interval == "step"
|
|
lightning_model.log(
|
|
f"train{path}/grad_norm",
|
|
grad_norm_val,
|
|
on_step=on_step,
|
|
on_epoch=not on_step,
|
|
)
|
|
|
