GradientNormLogger#

class GradientNormLogger(norm_type=2.0, log_per_layer=False, log_frequency=1, prefix='grad_norm')[source]#

Bases: HookBase

Log model gradient norms to the writer after each training step.

Runs in after_step every log_frequency steps: computes the aggregate gradient norm over all model parameters (for the configured norm_type) and writes it under {prefix}/total. With log_per_layer=True it additionally writes each named parameter’s gradient norm under {prefix}/layers/<name>. Registered as GradientNormLogger.

Parameters:

  • norm_type (float) – Order of the norm to compute, e.g. 2.0 for L2 or float("inf") for max-abs. Defaults to 2.0.

  • log_per_layer (bool) – If True, also log per-parameter gradient norms (verbose). Defaults to False.

  • log_frequency (int) – Compute and log every this many steps. Defaults to 1.

  • prefix (str) – Namespace prefix for the writer keys. Defaults to "grad_norm".

Note

No-ops when trainer.writer is absent/None or does not expose add_scalar. Reads gradients after backward, so it reflects the un-clipped gradients of that step.

Example

Add to cfg.hooks; after every log_frequency steps it writes the total gradient norm to the experiment writer (W&B/TensorBoard):

hooks = [dict(type="GradientNormLogger", log_frequency=50)]
# → logs scalar  "grad_norm/total"  every 50 optimizer steps
#   (with log_per_layer=True, also "grad_norm/layers/<name>" per param)

The norm helper is pure and can be exercised standalone:

>>> import torch
>>> from pimm.engines.hooks.diagnostics import GradientNormLogger
>>> p = torch.nn.Parameter(torch.tensor([3.0, 4.0]))
>>> (0.5 * (p ** 2).sum()).backward()   # grad = [3., 4.]
>>> float(GradientNormLogger()._compute_grad_norm([p]))
5.0
after_step()[source]#

Log gradient norms after each training step.