torch-surgeon

Real-time gradient pathology detection for PyTorch — in 2 lines.

A loss curve is a lagging indicator. By the time it shows a problem, vanishing or exploding gradients have been compounding for hundreds of steps. torch-surgeon attaches diagnostic hooks to your model and surfaces per-layer pathologies in real time, before they compound into an unrecoverable run.

Install

pip install torch-surgeon

Usage

from torch_surgeon import Surgeon

surgeon = Surgeon(model, rules="default")
surgeon.attach()

# ... your existing training loop, unchanged ...
for epoch in range(epochs):
    loss = criterion(model(x), y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

report = surgeon.report()   # per-layer stats dict
surgeon.detach()            # clean removal of all hooks

What it detects

Custom rules

surgeon = Surgeon(model, rules={
    "vanishing_threshold": 0.01,   # default
    "exploding_threshold": 100.0,  # default
    "stagnant_steps": 50,          # default
    "log_every": 10,               # print summary every N steps
    "plot": True,                  # live matplotlib plot
    "verbose": True,
})

How it works

torch-surgeon uses PyTorch's register_full_backward_hook API to intercept gradients at every leaf layer during the backward pass. Statistics (mean, std, norm) are computed inside the hook and the raw gradient tensor is discarded immediately — keeping overhead under 1% on typical training loops.

Pathology detection uses an exponential moving average (EMA) baseline per layer rather than fixed thresholds — so it generalises across architectures without manual tuning.

Performance

Sub-1% training overhead on standard loops. Validated against 100-step timing benchmarks on Linear/ReLU networks. Stats computed in-hook; no tensors stored between steps.

License

MIT