log_gradient_norm

class LogGradientNorm

Log total and per-parameter gradient norm at every training step.

on_before_optimizer_step(trainer: Trainer, pl_module: LightningModule, optimizer: Optimizer) → None

Log the gradients before the optimizer step.

class LogParameterNorm

Log total and per-parameter norm at every training step.

on_before_optimizer_step(trainer: Trainer, pl_module: LightningModule, optimizer: Optimizer) → None

Log the parameter norms before the optimizer step.

parameter_norm(module: Module, norm_type: float | int | str, group_separator: str = '/', learnable_only=True) → Dict[str, float]

Compute each parameter’s norm and their overall norm.

The overall norm is computed over all parameters together, as if they were concatenated into a single vector.

Based on lightning.pytorch.utilities.grad_norm.

Parameters:

• module – 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 parameter norms in their own subfolder instead of the logs one.

• learnable_only – Whether to only consider parameters that have a gradient.

Returns:

The dictionary of p-norms of each parameter and

a special entry for the total p-norm of the parameters viewed as a single vector.

Return type:

norms