Make the gradient-norm probe safe to leave on

No architectural change was involved in getting the gradient norm -- the model,
loss, optimiser, data pipeline and resolved training config are untouched, and
the logging attaches through the public add_callback API. But the probe patches
torch.nn.utils.clip_grad_norm_, and reviewing that for what goes wrong in
practice turned up three things worth fixing.

Ultralytics fires on_train_end only on success, so a crash left the patch
installed and the second model's probe then wrapped the wrapper, making
restoration hand back a wrapper instead of the real function. install() now
marks the function it creates and refuses to wrap a marked one, and train.py
calls close() from a finally block.

float() on a CUDA tensor synchronises the device, and clip_grad_norm_ does not
otherwise sync since error_if_nonfinite defaults to False. Converting every step
added a stall to every step of the run. Norms are kept as detached tensors and
converted once per epoch instead.

Gradient norm is now its own flag, --log-grad-norm, so loss and metric curves no
longer require patching anything in torch.

Fixes a NameError the refactor introduced: _run_training referenced model_name
from the enclosing scope.

tests/test_grad_norm_probe.py pins the behaviour against a stub torch, needing
neither a GPU nor a 2 GB install: restoration, refusal to nest, pass-through of
the caller's value, and that an uninstalled probe leaves torch alone.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

README.md

@@ -419,18 +419,33 @@ the process is changed.
 
 ### Getting at the gradient norm
 
+This is the one part of the logging that is not a callback, and it is opt-in
+separately as `--log-grad-norm`. Losses and metrics come from callbacks that only
+read `trainer.*`; without this flag nothing in torch is touched.
+
 `BaseTrainer.optimizer_step` calls `clip_grad_norm_`, whose return value is the
 total pre-clip gradient norm, and discards it. It then calls `zero_grad()` in the
 same method, and **no callback fires between the two** — the nearest,
 `on_train_batch_end`, runs after the gradients are already gone, so it cannot be
-recomputed either.
-
-Wrapping `clip_grad_norm_` for the duration of training is the only hook
-available. It is also the most accurate one: the captured value is taken after
-`scaler.unscale_()`, so it is a true unscaled norm rather than an AMP-scaled one,
-and it is the same number the optimiser acted on. Summarised per epoch rather
-than per step, since per-step would be thousands of points that do not line up
-with the losses.
+recomputed either. Wrapping `clip_grad_norm_` for the duration of training is the
+only hook available, and the most accurate one: the value is captured after
+`scaler.unscale_()`, so it is a true unscaled norm and the same number the
+optimiser acted on.
+
+The wrapper only reads a return value. It never touches gradients, the optimiser
+or the loss, so it cannot change what the model learns. Three things it could
+still have cost you, all handled:
+
+| Risk | Handling |
+| --- | --- |
+| `on_train_end` does not fire when training raises, so the patch survives and the second model's probe wraps the wrapper | `install()` marks the function it creates and refuses to wrap a marked one; `train.py` calls `close()` from a `finally` |
+| `float()` on a CUDA tensor synchronises the device, and `clip_grad_norm_` does not otherwise sync (`error_if_nonfinite` is False), so converting per step would stall every step | norms are kept as detached tensors and converted once per epoch |
+| Coupling it to `--wandb` would force the patch on anyone wanting loss curves | separate flag |
+
+[`tests/test_grad_norm_probe.py`](tests/test_grad_norm_probe.py) pins all of it
+against a stub torch — no GPU, no 2 GB install — covering restoration, refusal to
+nest, pass-through of the caller's value, and that constructing a probe without
+installing it leaves torch untouched.
 
 `grad_norm/clipped_fraction` is the series to read first: the share of steps
 whose norm exceeded the clip of 10. Near 1.0 during warmup is normal. If it stays

notebooks/colab_train.ipynb

@@ -139,7 +139,7 @@
     "    --data /content/data/data.yaml \\\n",
     "    --models yolo11n.pt yolo11s.pt \\\n",
     "    --epochs 100 --imgsz 640 --batch 16 --device 0 \\\n",
-    "    --wandb --wandb-project cone-distance \\\n",
+    "    --wandb --wandb-project cone-distance --log-grad-norm \\\n",
     "    --project /content/runs"
    ]
   },
@@ -163,7 +163,12 @@
     "gradients at norm 10, and this is the share of steps that hit the clip. It sits\n",
     "near 1.0 during warmup, which is normal; if it stays there once warmup ends, the\n",
     "effective step size is being set by the clip rather than by `lr0`, and lowering\n",
-    "the learning rate will do more than tuning anything else."
+    "the learning rate will do more than tuning anything else.\n",
+    "\n",
+    "Gradient norm is the one thing here that patches a torch function, because\n",
+    "Ultralytics discards the value and fires no callback while the gradients are\n",
+    "live. Dropping `--log-grad-norm` above removes that patch entirely and keeps\n",
+    "every loss and metric curve."
    ]
   },
   {

src/train/train.py

@@ -118,17 +118,36 @@ def train_one(model_name: str, data_yaml: Path, args) -> None:
 
     model = YOLO(model_name)
 
+    close_logger = None
     if args.wandb:
         from src.train import wandb_logger
 
-        wandb_logger.attach(
+        close_logger = wandb_logger.attach(
             model,
             project=args.wandb_project,
             run_name=f"{Path(model_name).stem}-{args.epochs}ep",
             extra_config={"dataset": str(data_yaml)},
+            log_grad_norm=args.log_grad_norm,
         )
-        print(f"logging to W&B project {args.wandb_project!r}")
+        print(f"logging to W&B project {args.wandb_project!r}"
+              + (" (with gradient norm)" if args.log_grad_norm else ""))
 
+    try:
+        _run_training(model, model_name, data_yaml, args)
+    finally:
+        # Ultralytics fires on_train_end only on success, and the gradient-norm
+        # probe patches a torch function that must be put back either way --
+        # otherwise a crash training the first model leaves a wrapper in place
+        # for the second.
+        if close_logger is not None:
+            close_logger()
+
+    metrics = model.val(data=str(data_yaml), imgsz=args.imgsz, device=args.device)
+    print_per_class_metrics(metrics, model.names)
+    print(f"\nweights: {args.project}/{Path(model_name).stem}/weights/best.pt")
+
+
+def _run_training(model, model_name: str, data_yaml: Path, args) -> None:
     model.train(
         data=str(data_yaml),
         epochs=args.epochs,
@@ -142,10 +161,6 @@ def train_one(model_name: str, data_yaml: Path, args) -> None:
         seed=args.seed,
     )
 
-    metrics = model.val(data=str(data_yaml), imgsz=args.imgsz, device=args.device)
-    print_per_class_metrics(metrics, model.names)
-    print(f"\nweights: {args.project}/{Path(model_name).stem}/weights/best.pt")
-
 
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__,
@@ -167,6 +182,12 @@ def main() -> None:
                         help="log train/val loss and gradient norm to Weights & Biases; "
                              "authenticate with $WANDB_API_KEY")
     parser.add_argument("--wandb-project", default="cone-distance")
+    parser.add_argument("--log-grad-norm", action="store_true",
+                        help="also log gradient norm. Needs --wandb. Unlike the loss "
+                             "callbacks, this patches torch.nn.utils.clip_grad_norm_ "
+                             "for the duration of training, because Ultralytics "
+                             "discards the value and fires no callback while the "
+                             "gradients are still live.")
     args = parser.parse_args()
 
     if args.wandb and not os.environ.get("WANDB_API_KEY"):

src/train/wandb_logger.py

@@ -29,6 +29,11 @@ import numpy as np
 CLIP_MAX_NORM = 10.0
 
 
+# Set on the patched function so a second install can recognise its own work and
+# refuse to wrap a wrapper.
+_PATCH_MARKER = "_cone_distance_grad_norm_probe"
+
+
 class _GradNormProbe:
     """Capture the gradient norm that Ultralytics computes and throws away.
 
@@ -43,24 +48,45 @@ class _GradNormProbe:
     only hook available. It is also the most accurate one: the value captured
     is taken after `scaler.unscale_()`, so it is a true unscaled norm rather
     than an AMP-scaled one, and it is the same number the optimiser acted on.
+
+    The wrapper only reads a return value -- it never touches gradients, the
+    optimiser or the loss -- so it cannot change what the model learns. Two
+    things it could still cost you, both handled here:
+
+      * `float()` on a CUDA tensor forces a device synchronisation, and
+        `clip_grad_norm_` does not otherwise sync (`error_if_nonfinite` is False
+        by default). Doing that every step would add a stall to every step of
+        the run, so norms are kept as detached tensors and converted once per
+        epoch instead.
+      * If training raises, `on_train_end` never fires. `install()` is therefore
+        idempotent, and `train.py` removes the probe from a `finally` block, so
+        a crash in the first model cannot leave a wrapper in place for the
+        second.
     """
 
     def __init__(self) -> None:
         self._original = None
-        self.norms: list[float] = []
+        self._norms: list = []
 
     def install(self) -> None:
         import torch
 
-        if self._original is not None:
+        current = torch.nn.utils.clip_grad_norm_
+        if self._original is not None or getattr(current, _PATCH_MARKER, False):
+            # Already patched, by us or by a probe whose removal was skipped.
+            # Wrapping again would nest wrappers and make restoration wrong.
             return
-        self._original = torch.nn.utils.clip_grad_norm_
+        self._original = current
 
         def recording_clip_grad_norm(parameters, max_norm, *args, **kwargs):
             total_norm = self._original(parameters, max_norm, *args, **kwargs)
-            self.norms.append(float(total_norm))
+            # detach() only, no float(): converting here would synchronise the
+            # device on every optimiser step.
+            self._norms.append(total_norm.detach() if hasattr(total_norm, "detach")
+                               else total_norm)
             return total_norm
 
+        setattr(recording_clip_grad_norm, _PATCH_MARKER, True)
         torch.nn.utils.clip_grad_norm_ = recording_clip_grad_norm
 
     def remove(self) -> None:
@@ -69,13 +95,15 @@ class _GradNormProbe:
         if self._original is not None:
             torch.nn.utils.clip_grad_norm_ = self._original
             self._original = None
+        self._norms.clear()
 
     def summarise(self) -> dict[str, float]:
         """Per-epoch summary. Logging every step would be thousands of points
-        per epoch and would not line up with the losses."""
-        if not self.norms:
+        per epoch and would not line up with the losses. This is also the only
+        place the device is synchronised -- once per epoch, not once per step."""
+        if not self._norms:
             return {}
-        norms = np.array(self.norms)
+        norms = np.array([float(n) for n in self._norms])
         return {
             "grad_norm/mean": float(norms.mean()),
             "grad_norm/max": float(norms.max()),
@@ -86,18 +114,28 @@ class _GradNormProbe:
         }
 
     def reset(self) -> None:
-        self.norms.clear()
+        self._norms.clear()
 
 
-def attach(model, project: str, run_name: str, extra_config: dict | None = None) -> None:
+def attach(model, project: str, run_name: str, extra_config: dict | None = None,
+           log_grad_norm: bool = False):
     """Register W&B callbacks on one Ultralytics model.
 
     Call once per model. Each model gets its own W&B run, so the YOLO11n and
     YOLO11s arms appear as two comparable runs rather than one interleaved mess.
+
+    Returns a `close()` callable. Call it from a `finally` block: Ultralytics
+    only fires `on_train_end` on success, and with `log_grad_norm` the probe
+    patches a torch function that must be put back even when training raises.
+
+    `log_grad_norm` is separate from W&B logging on purpose. Losses and metrics
+    come from callbacks that only read `trainer.*`; the gradient norm needs a
+    patched `torch.nn.utils.clip_grad_norm_`. Leaving it off means nothing in
+    torch is touched.
     """
     import wandb
 
-    probe = _GradNormProbe()
+    probe = _GradNormProbe() if log_grad_norm else None
 
     def on_pretrain_routine_start(trainer):
         wandb.init(
@@ -106,10 +144,12 @@ def attach(model, project: str, run_name: str, extra_config: dict | None = None)
             config={**vars(trainer.args), **(extra_config or {})},
             reinit=True,
         )
-        probe.install()
+        if probe is not None:
+            probe.install()
 
     def on_train_epoch_start(trainer):
-        probe.reset()
+        if probe is not None:
+            probe.reset()
 
     def on_train_epoch_end(trainer):
         step = trainer.epoch + 1
@@ -122,13 +162,20 @@ def attach(model, project: str, run_name: str, extra_config: dict | None = None)
         # step as the training losses above so the two curves overlay.
         step = trainer.epoch + 1
         wandb.log(trainer.metrics, step=step)
-        summary = probe.summarise()
-        if summary:
-            wandb.log(summary, step=step)
+        if probe is not None:
+            summary = probe.summarise()
+            if summary:
+                wandb.log(summary, step=step)
 
     def on_train_end(trainer):
-        probe.remove()
-        wandb.finish()
+        close()
+
+    def close():
+        """Restore torch and end the run. Safe to call more than once."""
+        if probe is not None:
+            probe.remove()
+        if wandb.run is not None:
+            wandb.finish()
 
     for event, handler in (
         ("on_pretrain_routine_start", on_pretrain_routine_start),
@@ -138,3 +185,5 @@ def attach(model, project: str, run_name: str, extra_config: dict | None = None)
         ("on_train_end", on_train_end),
     ):
         model.add_callback(event, handler)
+
+    return close

tests/test_grad_norm_probe.py

@@ -0,0 +1,97 @@
+"""The gradient-norm probe patches a torch function, so it gets its own test.
+
+Runs against a stub torch, so it needs neither a GPU nor a 2 GB install. What
+it pins down is the behaviour that would be expensive to discover during a real
+100-epoch run: that the patch restores cleanly, that it refuses to wrap itself,
+and that it never alters the value the optimiser receives.
+
+    python -m tests.test_grad_norm_probe
+"""
+
+from __future__ import annotations
+
+import sys
+import types
+
+
+def install_stub_torch() -> types.ModuleType:
+    """Minimal stand-in exposing only torch.nn.utils.clip_grad_norm_."""
+    torch = types.ModuleType("torch")
+    nn = types.ModuleType("torch.nn")
+    utils = types.ModuleType("torch.nn.utils")
+
+    class FakeNorm(float):
+        """Stands in for the tensor clip_grad_norm_ returns."""
+        def detach(self):
+            return self
+
+    calls = []
+
+    def clip_grad_norm_(parameters, max_norm, *args, **kwargs):
+        calls.append(max_norm)
+        return FakeNorm(12.5)
+
+    clip_grad_norm_.calls = calls
+    utils.clip_grad_norm_ = clip_grad_norm_
+    nn.utils = utils
+    torch.nn = nn
+    sys.modules["torch"] = torch
+    sys.modules["torch.nn"] = nn
+    sys.modules["torch.nn.utils"] = utils
+    return torch
+
+
+def main() -> None:
+    torch = install_stub_torch()
+    from src.train.wandb_logger import _GradNormProbe
+
+    pristine = torch.nn.utils.clip_grad_norm_
+
+    # 1. Install, and the value the caller gets back is untouched.
+    a = _GradNormProbe()
+    a.install()
+    assert torch.nn.utils.clip_grad_norm_ is not pristine, "probe did not patch"
+    returned = torch.nn.utils.clip_grad_norm_([], 10.0)
+    assert float(returned) == 12.5, returned
+    assert pristine.calls == [10.0], "the real function was not called through"
+
+    # 2. A second probe must not wrap the wrapper. This is the case that arises
+    #    when training the first model raises and its probe is never removed.
+    b = _GradNormProbe()
+    b.install()
+    assert torch.nn.utils.clip_grad_norm_ is not pristine
+    depth_marker = torch.nn.utils.clip_grad_norm_
+    b.install()
+    assert torch.nn.utils.clip_grad_norm_ is depth_marker, "probe nested on itself"
+
+    # 3. Removing the probe that actually patched restores the real function.
+    b.remove()   # b never owned the patch, so this is a no-op
+    assert torch.nn.utils.clip_grad_norm_ is not pristine, "non-owner removed the patch"
+    a.remove()
+    assert torch.nn.utils.clip_grad_norm_ is pristine, "patch was not restored"
+    a.remove()   # idempotent
+
+    # 4. Summaries, including the clipped fraction against the clip of 10.
+    c = _GradNormProbe()
+    c.install()
+    for _ in range(4):
+        torch.nn.utils.clip_grad_norm_([], 10.0)      # each returns 12.5 > 10
+    summary = c.summarise()
+    assert summary["grad_norm/mean"] == 12.5, summary
+    assert summary["grad_norm/max"] == 12.5, summary
+    assert summary["grad_norm/clipped_fraction"] == 1.0, summary
+    c.reset()
+    assert c.summarise() == {}, "reset did not clear the epoch"
+    c.remove()
+    assert torch.nn.utils.clip_grad_norm_ is pristine
+
+    # 5. The default path must not touch torch at all.
+    before = torch.nn.utils.clip_grad_norm_
+    _GradNormProbe()          # constructed but never installed
+    assert torch.nn.utils.clip_grad_norm_ is before
+
+    print("GRAD NORM PROBE TEST PASSED")
+
+
+if __name__ == "__main__":
+    main()