hiddenstateresearch/visualtorch

VisualTorch – Quick Overview

What it is – VisualTorch is a Python library that turns a PyTorch model into a readable diagram. It works by running a dummy forward pass, tracing the operations, and then drawing the architecture in one of three styles:

• Flow style – layers stacked left‑to‑right, similar to a data‑flow chart.
• Graph style – a more traditional node‑edge graph showing connections.
• LeNet style – a compact, textbook‑like layout.

The library can also export an animated GIF that reveals the model layer‑by‑layer.

---

Why it matters for AI practitioners

When you build or share a neural network, a clear visual representation helps:

• Communication – explain model structure to collaborators or reviewers.
• Debugging – spot missing connections or unexpected branches.
• Documentation – embed diagrams in papers, notebooks, or READMEs.

VisualTorch fills a niche between simple text summaries (torchsummary) and heavyweight graph tools (torchviz). It adds a unified tracing backend that works for both nn.Sequential and arbitrary custom modules.

---

Core Features (as described in the README)

Feature	What you get
Three rendering styles	flow, graph, and lenet layouts, each with configurable colors and fonts.
Animated GIF export	visualtorch.animate(model, input_shape, style=…) creates a step‑by‑step reveal of the architecture.
Dimension labeling	Option show_dimension=True adds the shape of every output tensor to node labels.
MCP server integration	Optional visualtorch[mcp] installs a generic stdio‑based server that can be called from any client to generate PNGs or GIFs without embedding the library.
Research‑ready citation	A DOI and BibTeX entry are provided for academic use.
Extensive docs & examples	Online docs host usage snippets, API reference, and a showcase of papers that have used the tool.
PyPI distribution	Install with pip install visualtorch (or with the [mcp] extra).

---

How to get started (from the README)

# Basic install
pip install visualtorch

# Install with the optional MCP server
pip install "visualtorch[mcp]"

import torch, visualtorch as vt

# Example model (any nn.Module works)
class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 16, 3)
        self.fc   = torch.nn.Linear(16*26*26, 10)
    def forward(self, x):
        x = torch.relu(self.conv(x))
        return self.fc(x.view(x.size(0), -1))

model = Net()
# Generate a static PNG in flow style
vt.plot(model, input_shape=(1, 1, 28, 28), style="flow", show_dimension=True)

# Create an animated GIF that reveals the model step‑by‑step
vt.animate(model, input_shape=(1, 1, 28, 28), style="graph")

---

Known limitations (explicitly called out)

• Tracing‑only – VisualTorch runs a real forward pass with a dummy input. Any data‑dependent control flow will only display the branch taken by that dummy input.
• Multi‑output layers – Nodes are sized based on the first output tensor; however, show_dimension=True still lists shapes for all outputs, and connections remain correct.
• No symbolic execution – The library cannot infer shapes for branches that depend on runtime tensor values.

---

Ecosystem fit & Related Tools

VisualTorch builds on ideas from:

• visualkeras – early inspiration for style and API.
• pytorchviz, torchview, pytorch‑summary – earlier tracing/summary utilities.

It differentiates itself by offering multiple ready‑made visual styles, animated reveals, and a client‑neutral MCP server that can be called from any language or toolchain.

---

Community & Maintenance

• Open‑source MIT license – free for commercial and academic use.
• Citation encouraged – a JOSS paper (2024) provides a DOI for academic credit.
• Contribution welcome – guidelines and a CONTRIBUTING.md are provided.
• Release process – governed by a GOVERNANCE.md file; version 1.0+ introduced breaking API changes but adds richer features.
• Usage in published research – the README lists papers in Nature, IEEE, MDPI that have employed VisualTorch.

---

Bottom line

VisualTorch is a purpose‑built, well‑documented library for turning PyTorch models into clean, customizable diagrams (static PNGs or animated GIFs). It is especially useful for researchers and engineers who need to communicate model architectures in papers, presentations, or documentation without writing custom plotting code.