torchview Documentation

The documentation structure follows a standard MkDocs layout:

• mkdocs.yml: The main configuration file for the documentation site.
• docs/: The directory containing all documentation content.
  • index.md: The homepage of the documentation site.
  • Other .md files, images, and assets used within the documentation.

The visual graphs produced by torchview consist of nodes and directed edges. Nodes are categorized into three types:

1. Tensor Node (Bright Yellow): Labeled as {tensor-name}{depth}: {tensor-shape}. Names include input-tensor, hidden-tensor, or output-tensor.
2. Function Node (Bright Blue): Labeled as {Function-name}{depth}: {input and output shape}.
3. Module Node (Bright Green): Labeled as {Module-name}{depth}: {input and output shape}.

When roll=True is used for recursive modules, numbers on edges indicate how many times that specific edge occurs during the forward pass (e.g., a '4' indicates the connection is used 4 times).

The rolling mechanism determines whether recursive computations are displayed as unique nodes or collapsed into shared nodes in the visual graph. This is controlled by the Roll parameter.

• Roll=False: Every computation step is displayed uniquely on the graph, even if the same module or function is used recursively. Each node is assigned a unique node_id based on its Python object ID (e.g., id(tensor_node_1)).
• Roll=True: Recursively used modules are identified as the same node on the visual graph to simplify the visualization. The identification logic depends on the node type (see Node Identification Logic).

When using VSCode, the default SVG format may not render as expected. To ensure correct visualization in Jupyter notebooks within VSCode, set the graphviz format to 'png'.

import graphviz
# when running on VSCode run the below command
# svg format on vscode does not give desired result
graphviz.set_jupyter_format('png')

The torchview Model Gallery provides a collection of SVG vector visualizations for various deep learning architectures across different domains (Computer Vision, NLP, Recommender Systems, etc.).

Because the visualizations are in SVG vector format, you can:

• Click on any image to expand it.
• Zoom in using your browser to inspect fine-grained details of the model architecture (e.g., skip connections, attention heads, or bottleneck structures).

To add a new model architecture to the gallery, follow these requirements:

1. Ensure the model architecture is representative of its class.
2. Use torchview to generate the visualization in SVG format to maintain clarity and scalability.
3. Submit a Pull Request (PR) to the repository.

To use torchview, you need to install the package itself along with graphviz for rendering the graphs.

! pip install -q torchview
! pip install -q -U graphviz

#! pip install -q torchview
! pip install -q -U graphviz

The TorchView documentation is built using MkDocs. You can use the following commands to manage the documentation site:

• Create a new project: mkdocs new [dir-name]
• Start a live-reloading server for local preview: mkdocs serve
• Build the static documentation site: mkdocs build
• View help information: mkdocs -h

To use torchview, you must use a PyTorch version greater than 1.6. This is because torch.Tensor did not have __torch_function__ as a class method in versions 1.6 and earlier, which is required for the package to function.

Additionally, be aware of the following behavior for specific operators in older PyTorch versions:

• PyTorch 1.7.1, 1.8, and 1.9: F.linear and F.embedding return NotImplemented when called via __torch_function__ on a subclass. torchview includes internal workarounds to handle these specific versions and ensure they return RecorderTensor instead of standard torch.Tensor.