visualtorch

visualtorch

MCP.Pizza Chef: willyfh

Point it at the code for a PyTorch model and it draws the layers as a labeled image, in a flow, graph, or classic LeNet style. It can also produce an animated version that reveals the model one layer at a time, which is handy for a paper figure, a slide, or a lecture. Everything runs on your own machine with no account or key. You install it alongside Python and connect it to an assistant that supports MCP.

Coding
Images/Design

Use This MCP server To

Draw a diagram of my PyTorch model for a paper Make an animated reveal of a model's layers for a talk Show the shape of the data at each layer Save a model diagram as an image file Produce a figure without leaving my assistant

README

VisualTorch

python pytorch Downloads Run Tests Documentation Status

VisualTorch aims to help visualize Torch-based neural network architectures. It currently supports generating flow-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. Its original visual styles were inspired by visualkeras, pytorchviz, pytorch-summary, and torchview; since then, it has grown its own unified tracing backend and architecture-handling logic well beyond its origins.

Note: 1.0+ is a major release with breaking API changes, but with significantly better features and algorithms - upgrading is recommended. For the old API, use 0.2.5 or older.

Limitation: VisualTorch traces a real forward pass to build the diagram, which has an inherent limitation shared by any tracing-based approach (not a bug, and not fixable without full symbolic execution): models with data-dependent control flow (e.g. a branch only taken if a tensor value crosses some threshold) only show whichever branch the traced dummy input happened to take. Separately, a layer that returns multiple meaningful output tensors (e.g. a custom multi-task head, or nn.LSTM's (output, (h_n, c_n))) still has its node's size based on only its first tensor; with show_dimension=True, every output tensor's shape is shown in the label, not just the first. Downstream connections are correct either way. Contributions are welcome!

VisualTorch Examples

Animated Reveal

Every style can also render as an animated GIF, revealing the model one layer/column at a time, via visualtorch.animate(model, input_shape, style=...) - see it in action for flow, graph, and lenet styles.

Animated VisualTorch Example

visualtorch FAQ

Which apps does it work in?
Any assistant that supports MCP over a local connection. No special plugin or extension is needed.
Do I need a key or an account?
No. It runs locally and only needs Python and PyTorch on your machine.
Can I use this to make a figure for a research paper?
Yes — it produces publication-style diagrams and has been used in published research.
What kinds of models does it handle?
PyTorch models, both simple stacked ones and custom ones you have written yourself.
How hard is the setup?
You install one Python package with its optional MCP piece, then point your assistant at it.
Are the diagrams always exact?
Almost always — a model that branches based on the data it sees will only show the path actually taken.