# gradviz

[![PyPI version](https://badge.fury.io/py/gradviz.svg)](https://badge.fury.io/py/gradviz)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

**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

```bash
pip install gradviz
```

## ⚡ Quickstart

```python
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
![Alt text](heatmap.png)

Heatmap of gradient norms by layer
![Alt text](lines.png)

## 🔧 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

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

## API Reference

```python
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.
