Visualize and debug PyTorch model tensors
Project description
🧠 TensorViz
Visualize, debug, and inspect your PyTorch models like never before.
TensorViz is an AI-powered Python package that automatically logs and visualizes the internals of PyTorch models during forward/backward passes — including input/output shapes, data types, gradients, and more. It generates an interactive HTML report with collapsible groups, Graphviz diagrams, and a clean structure for inspecting every layer.
🚀 Features
- ✅ Auto-log layer name, input/output shape, data type
- ✅ Optional gradient stats (mean ± std) after
.backward() - ✅ Collapsible
<details>sections for layer groups - ✅ Inline SVG diagram via Graphviz
- ✅ Lightweight, minimal dependencies
- ✅ Python API + CLI
- ✅ Works with all
torch.nn.Modules
📦 Installation
pip install tensorviz
You’ll also need Graphviz for the SVG model diagram:
# Ubuntu/Debian
sudo apt install graphviz
# macOS (via Homebrew)
brew install graphviz
# Optional Python binding
pip install graphviz
✨ Quick Start
Python API
import torch
import torchvision.models as models
from tensorviz import visualize_model
model = models.resnet18()
input_tensor = torch.randn(1, 3, 224, 224)
visualize_model(model, input_tensor, save_dir="tensorviz_report")
CLI
tensorviz report model.pt --input-shape 1 3 224 224
This generates:
tensorviz_report/logs.jsontensorviz_report/graph.svgtensorviz_report/report.html
🧰 Usage Options
visualize_model(
model,
input_tensor=torch.randn(1, 3, 224, 224),
output_file="tensorviz_report/report.html",
include_grad=True, # Capture gradient statistics
include_graph=True, # Include Graphviz model diagram
html_template=None # Optional: custom HTML template path
)
🔍 How It Works
TensorViz is built on:
- PyTorch forward hooks for real-time logging
- Backward hooks for optional
.gradcapture - Jinja2 for HTML templating
- Graphviz via
torchviz(or manual tracing) - Native HTML5 tags (
<details>,<summary>) for collapsibility
📁 Project Structure
tensorviz/
├── __init__.py
├── hooks.py # Layer hooks (forward/backward)
├── render.py # HTML generation logic
├── svg_generator.py # Graphviz export logic
├── cli.py # CLI entry point
├── templates/
│ └── report.html.j2 # HTML report template
├── static/ # Optional CSS/assets
tests/
examples/
setup.py
README.md
LICENSE
📜 License
TensorViz is released under the MIT License.
MIT License
Copyright (c) 2025 Divyansh Saraswat
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the “Software”), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
...
See LICENSE for the full legal terms.
👇 Optional License Alternatives
If your organization or project requires a different license, the following alternatives are compatible:
- Apache License 2.0
- BSD 3-Clause License
- GNU GPL v3.0 (strict copyleft)
Please fork and adapt the license header accordingly if needed.
👥 Contributing
Contributions are welcome!
- Fork the repository
- Create a feature branch
- Make your changes
- Run tests (
pytest) - Submit a pull request with a clear description
🧪 Testing
Run tests using:
pytest tests/
Or run a manual test:
python examples/run_resnet.py
🙋 FAQ
Q: Does it work with custom layers?
Yes. Any torch.nn.Module subclass will be traced automatically.
Q: Does it work with HuggingFace Transformers?
Yes. TensorViz supports torchvision, transformers, and custom networks.
Q: Is Graphviz optional?
Yes. The HTML report works without it, but the model diagram will be skipped.
✨ Acknowledgments
🔗 Related Projects
Made with ❤️ by Divyansh Saraswat
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file tensorviz-0.1.6.tar.gz.
File metadata
- Download URL: tensorviz-0.1.6.tar.gz
- Upload date:
- Size: 11.1 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.1.0 CPython/3.11.0
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
b706a3bed981fa0cf3cc9366f5907c00fd2bd6a3e866e0c7be55cd5a608daf54
|
|
| MD5 |
3aef8b361ef8971a16cdd4a288a6b4d9
|
|
| BLAKE2b-256 |
4c48c286ce4e7d499a21c08f03cca54217b2397c9e545f4bb2debaf58d51464e
|
File details
Details for the file tensorviz-0.1.6-py3-none-any.whl.
File metadata
- Download URL: tensorviz-0.1.6-py3-none-any.whl
- Upload date:
- Size: 8.5 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.1.0 CPython/3.11.0
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
5c9b8e49df39a6b800b734908823cf919aa778d6af30800b7349e64d0a5efd5d
|
|
| MD5 |
5d262dad8380aeda2d336ed3fb3c0320
|
|
| BLAKE2b-256 |
b3bd5fd136eadfccec1298f1e49ddadc24841e21972c89c0c218f54250f59954
|