Torchviz

0.0.3 · active · verified Tue Apr 14

Torchviz is a small Python package designed to create visualizations of PyTorch execution graphs and traces. It provides a way to visually inspect the computational flow of a neural network model, which is helpful for understanding architecture, debugging, and optimization. The current version is 0.0.3, with releases typically tied to PyTorch ecosystem updates.

Warnings

breaking Torchviz critically depends on the external Graphviz software package being installed on your operating system, not just the Python `graphviz` package. Without a system-wide Graphviz installation (e.g., `brew install graphviz` on macOS, `sudo apt-get install graphviz` on Ubuntu, or installing from graphviz.org for Windows), `make_dot` will not be able to render images, often failing silently or with obscure errors about missing executables.
Fix: Install Graphviz on your system using your operating system's package manager or by downloading from the official Graphviz website (graphviz.org).
gotcha For `make_dot` to generate a meaningful computational graph, at least one of the input tensors to the operation being visualized must have `requires_grad=True`. If no tensor requires a gradient, PyTorch's autograd engine (which `torchviz` inspects) won't track operations, resulting in an empty or incomplete graph.
Fix: Ensure that the input tensor used in the forward pass has `requires_grad=True` (e.g., `torch.randn(..., requires_grad=True)`).
deprecated The `make_dot_from_trace` function, intended for use with `torch.jit.trace`, is noted to be less robust and 'does not always work' according to the official documentation, especially with newer PyTorch versions. Past issues have reported it not working with PyTorch 1.0.
Fix: Prefer using `make_dot` with the output of a forward pass and `model.named_parameters()` for reliable graph visualization.
gotcha The `show_attrs=True` and `show_saved=True` parameters, which provide additional details about the graph nodes (attributes and saved tensors for the backward pass), are only supported with PyTorch versions 1.9 and later. Using them with older PyTorch versions might not have the intended effect or could lead to errors.
Fix: Ensure your PyTorch installation is version 1.9 or higher to fully utilize `show_attrs` and `show_saved`. If using an older version, omit these parameters.

Install

pip install torchviz Install Python package
pip install graphviz # System-level Graphviz is also required (e.g., `brew install graphviz` on macOS, `sudo apt-get install graphviz` on Ubuntu) Install Graphviz (Python package and system-level)

Imports

make_dot
```
from torchviz import make_dot
```
This is the primary function for visualizing computational graphs from a PyTorch Variable.
make_dot_from_trace
```
from torchviz import make_dot_from_trace
```
While available, `make_dot_from_trace` using `torch.jit.trace` is noted to be less reliable and 'does not always work', especially with newer PyTorch versions. Prefer `make_dot` for general use.

Quickstart

This quickstart defines a simple PyTorch neural network, performs a forward pass with an input tensor that requires gradients, and then uses `make_dot` to generate and optionally save a visualization of the computational graph. Ensure Graphviz is installed on your system for rendering.

import torch
import torch.nn as nn
from torchviz import make_dot

# Define a simple neural network
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Instantiate the model
model = SimpleNN()

# Generate a random input, ensuring requires_grad=True for graph generation
input_data = torch.randn(1, 10, requires_grad=True)

# Perform a forward pass
output = model(input_data)

# Visualize the computational graph
graph = make_dot(output, params=dict(model.named_parameters()))

# To save the visualization to a file (e.g., PNG)
# graph.render("computational_graph", format="png", cleanup=True)

# In Jupyter/Colab, the graph can often be displayed directly by just calling it
# graph

view raw JSON →