gradviz 0.1.2

Lightweight gradient visualizer for PyTorch: record gradient norms and plot flows by layer/param.

gradviz

Lightweight gradient visualizer for PyTorch
Record gradient norms per parameter/layer during training, save them to CSV, and visualize gradient flow with line plots or heatmaps.

---

🚀 Why gradviz?

Training deep neural networks often runs into vanishing/exploding gradients problems. Debugging these issues usually means writing boilerplate hooks or printing raw tensors.
gradviz makes this simple:

• Attach once to your model
• Train as usual
• Save & visualize gradient norms across layers or parameters

Perfect for students, researchers, and anyone learning how backpropagation behaves.

---

📦 Installation

pip install gradviz

⚡ Quickstart

from gradviz import GradViz

gv = GradViz(model)
gv.attach()

for epoch in range(epochs):
    gv.set_epoch(epoch)
    for x, y in loader:
        opt.zero_grad()
        loss = model(x).loss(y)
        loss.backward()
        opt.step()
        gv.step()

gv.detach()
gv.save("gradviz.csv")
gv.plot(by="layer", topk=20)
gv.heatmap(by="layer", at_step=1000)

🖼️ Example Plots

Line plot of gradient norms over steps

Heatmap of gradient norms by layer

🔧 Features

• Record gradient L2 norms during training
• Works with any PyTorch model
• Save to CSV for later analysis
• Visualize:
  1. Line plots (by="layer" or by="param")
  2. Heatmaps at specific steps
• Lightweight (depends only on torch, numpy, pandas, matplotlib)
• Command-line interface (CLI) included

CLI Usage

gradviz plot gradviz.csv --by layer --topk 15
gradviz heatmap gradviz.csv --by param --step 500

API Reference

from gradviz import GradViz, GradVizConfig

1. GradViz(model, config=None) → main class
2. .attach() → hook gradients
3. .detach() → remove hooks
4. .step() → call after optimizer.step()
5. .set_epoch(epoch) → optionally record epoch index
6. .save(path) → save collected gradients to CSV
7. .plot(by="layer"|"param", topk=20) → line plots
8. .heatmap(by="layer"|"param", at_step=...) → heatmap

📂 Examples

See examples/demo_mnist.py for a full MNIST training demo.