pytorch-gradwatch 0.1.1.dev0
GradWatch: Advanced PyTorch Debugger with a Live TUI.
License: MIT
Requires Python: <3.13,>=3.8
Keywords: pytorch, debug, monitor, training, tui, rich, neural network, machine learning

<div align="center">
  <img src="https://raw.githubusercontent.com/im-mahdi-74/gradwatch/main/.github/assets/logo.png" alt="GradWatch Logo" width="200"/>
  <h1>GradWatch</h1>
  <p><b>Advanced PyTorch Debugger & Monitor with a Live TUI for MLOps.</b></p>
  <p>Never let NaNs ruin your training again. Proactively monitor gradients, activations, and neuron health in real-time to ensure robust deep learning training.</p>
  
  <!-- Badges -->
  <p>
    <a href="https://github.com/im-mahdi-74/gradwatch/actions/workflows/ci.yml"><img src="https://github.com/im-mahdi-74/gradwatch/actions/workflows/ci.yml/badge.svg" alt="CI Status"></a>
    <a href="https://pypi.org/project/gradwatch/"><img src="https://img.shields.io/pypi/v/gradwatch.svg" alt="PyPI version"></a>
    <a href="https://pypi.org/project/gradwatch/"><img src="https://img.shields.io/pypi/pyversions/gradwatch.svg" alt="Python versions"></a>
    <a href="https://codecov.io/gh/im-mahdi-74/gradwatch"><img src="https://codecov.io/gh/im-mahdi-74/gradwatch/branch/main/graph/badge.svg" alt="Coverage"></a>
    <a href="https://github.com/im-mahdi-74/gradwatch/blob/main/LICENSE"><img src="https://img.shields.io/pypi/l/gradwatch.svg" alt="License"></a>
  </p>
</div>

---

**GradWatch** is an advanced, production-ready PyTorch training debugger and real-time monitor. It's a simple, powerful context manager that wraps your training loop to provide:

-   💥 **Real-time Catastrophe Detection**: Proactively identifies and alerts for `NaN` or `Infinity` values in gradients, activations, or parameters, preventing silent training failures and saving valuable compute resources.
-   📊 **Live TUI Dashboard**: A beautiful and responsive terminal dashboard provides immediate visual feedback on your model's health during training.
-   🔎 **Post-Mortem Analysis**: Generate detailed reports from logged training data via the command-line interface (CLI) to diagnose issues after training, including exploding/vanishing gradients, dead neurons, and frozen parameters.
-   ✍️ **Persistent Logging**: All monitored metrics are logged to a SQLite database for historical analysis, and a CSV summary log tracks high-level run statistics.

## ✨ Features

-   **Gradient Monitoring**: Track L2 norm and max value of gradients to identify exploding or vanishing gradient problems.
-   **Dead Neuron Detection**: Monitor the percentage of "dead" (zero-output) neurons in `ReLU` activations.
-   **Parameter & Activation Tracking**: Keep an eye on the norms and max values of your model's parameters and activations.
-   **Notebook & Terminal Friendly**: Works seamlessly in both standard terminals (with a live TUI) and Jupyter/Colab notebooks (with simplified text updates).
-   **MLOps Ready**: Supports robust, reproducible training workflows with comprehensive logging and an adaptable API.
-   **Lightweight & Non-intrusive**: Utilizes PyTorch hooks and a Pythonic context manager API, adding minimal overhead to your training loop.

## 🚀 Quick Start

### 1. Installation

Install GradWatch using `pip`. Depending on your needs, you can install optional dependencies (extras):

```bash
# Minimal installation (core functionality)
pip install gradwatch

# To enable the live TUI dashboard
pip install "gradwatch[tui]"

# To enable the CLI plotting feature
pip install "gradwatch[cli]"

# For a full development environment (testing, linting, etc.)
pip install "gradwatch[dev]"

# For Notebook support (matplotlib, ipywidgets)
pip install "gradwatch[notebook]"

# To install everything
pip install "gradwatch[all]"
```

### 2. Usage

Wrap your training loop with the `gradwatch.watch` context manager. It's that simple!

```python
import torch
import torch.nn as nn
from gradwatch import watch
from gradwatch.config import GradWatchConfig, TrainingConfig, LoggingConfig

# 1. Define your model and optimizer
model = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 1))
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 2. Configure GradWatch
# Tell GradWatch about your training loop structure.
config = GradWatchConfig(
    training=TrainingConfig(
        total_epochs=10,
        steps_per_epoch=100,
    ),
    # Use the LoggingConfig object for logging settings
    logging=LoggingConfig(
        log_file='training_logs.db'
    )
)

# 3. Wrap your training loop with watch()
with watch(model, config=config) as gw:
    for epoch in range(10):
        for step in range(100):
            # Your standard training logic
            optimizer.zero_grad()
            inputs = torch.randn(64, 10)
            targets = torch.randn(64, 1)
            outputs = model(inputs)
            loss = (outputs - targets).pow(2).mean()
            loss.backward()
            optimizer.step()

            # 4. Update GradWatch with the latest loss
            gw.set_step(epoch, step, loss.item())

# If a NaN/Inf occurs, watch() will raise a CatastrophicTrainingError
# and stop the training immediately after logging the event.
```

## 🖥️ The TUI Dashboard

While your model trains, GradWatch displays a live dashboard in your terminal. After the run, the CLI provides a detailed summary:

```text
┌────────────────────────── Post-Mortem Analysis ───────────────────────────┐
│                📊 GradWatch Analysis Report for training.db                │
└───────────────────────────────────────────────────────────────────────────┘

1. Catastrophe Check:
✅ No NaN/Inf events were logged.

2. Gradient Health Report (Top 5 Offenders):
┌───────────────────────────────────────────────────────────────────────────┐
│                  ⚡ Top 5 Max Gradient Norms (Exploding Risk)                  │
├──────────────────┬──────────────────┤
│ Layer            │ Max Norm Value   │
├──────────────────┼──────────────────┤
│ layer2.weight    │ 2.34e+01         │
│ layer1.weight    │ 1.56e+01         │
└──────────────────┴──────────────────┘

3. Neuron Health Report (Top 5 Offenders):
┌───────────────────────────────────────────────────────────────────────────┐
│              🧠 Top 5 Layers with Most Dead Neurons (Avg %)               │
├──────────────────┬───────────────┤
│ Layer (ReLU)     │ Avg. Dead %   │
├──────────────────┼───────────────┤
│ relu1            │ 35.0%         │
└──────────────────┴───────────────┘
```

## CLI

After your training run, use the `gradwatch` CLI to perform a post-mortem analysis on the generated log file.

### Summary Report

Get a comprehensive health report for a completed run:

```bash
gradwatch summary training_logs.db
```

### Plot Metric History

Plot the history of a specific metric for a given layer:

```bash
gradwatch plot training_logs.db --layer layer1.weight --metric grad_norm
```

## 🤝 Contributing

Contributions are welcome! Please feel free to open an issue or submit a pull request.

1.  Fork the repository.
2.  Create a new branch (`git checkout -b feature/your-feature`).
3.  Install all dependencies for development: `pip install "gradwatch[all]"`
4.  Make your changes.
5.  Run tests, linting, and formatting checks: `ruff format . && ruff check . && mypy gradwatch/ && pytest`
6.  Commit your changes and push to your branch.
7.  Open a pull request.

## 📜 License

This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.

---

## 📝 To Do & Future Enhancements

-   [ ] Integration with MLOps platforms (e.g., MLflow, Weights & Biases)
-   [ ] GPU/Memory Usage Monitoring
-   [ ] Activation distribution drift detection
-   [ ] Customizable alert thresholds for gradients and dead neurons
-   [ ] Option to automatically save a model checkpoint on catastrophe detection

## 🧑‍💻 Contact

Developed by Mahdi Mosavi. You can find more of my work on [GitHub](https://github.com/im-mahdi-74).
Feel free to reach out via email: `mahdi.mosavi.nsa@gmail.com`.