Skip to main content

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.json
  • tensorviz_report/graph.svg
  • tensorviz_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 .grad capture
  • 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:

Please fork and adapt the license header accordingly if needed.


👥 Contributing

Contributions are welcome!

  1. Fork the repository
  2. Create a feature branch
  3. Make your changes
  4. Run tests (pytest)
  5. 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


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

tensorviz-0.1.6.tar.gz (11.1 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

tensorviz-0.1.6-py3-none-any.whl (8.5 kB view details)

Uploaded Python 3

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

Hashes for tensorviz-0.1.6.tar.gz
Algorithm Hash digest
SHA256 b706a3bed981fa0cf3cc9366f5907c00fd2bd6a3e866e0c7be55cd5a608daf54
MD5 3aef8b361ef8971a16cdd4a288a6b4d9
BLAKE2b-256 4c48c286ce4e7d499a21c08f03cca54217b2397c9e545f4bb2debaf58d51464e

See more details on using hashes here.

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

Hashes for tensorviz-0.1.6-py3-none-any.whl
Algorithm Hash digest
SHA256 5c9b8e49df39a6b800b734908823cf919aa778d6af30800b7349e64d0a5efd5d
MD5 5d262dad8380aeda2d336ed3fb3c0320
BLAKE2b-256 b3bd5fd136eadfccec1298f1e49ddadc24841e21972c89c0c218f54250f59954

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page