visualtorch 1.1.0

Architecture visualization of Torch models

🔥 VisualTorch 🔥

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. This tool is adapted from visualkeras, pytorchviz, and pytorch-summary.

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 two inherent limitations shared by any tracing-based approach (not bugs, and not fixable without full symbolic execution): (1) 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; (2) a layer that returns multiple meaningful output tensors (e.g. a custom multi-task head) only has its first tensor's shape reflected in that node's size/label - its downstream connections are still correct either way. Contributions are welcome!

Documentation

Online documentation is available at visualtorch.readthedocs.io.

The docs include usage examples, API references, and other useful information.

Installation

See the Installation page.

Used in Research

VisualTorch has been used in published research, including works published in Nature, IEEE, and MDPI.

See the Research Showcase page for the full list.

Examples

See the Usage Examples page.

Contributing

Please feel free to send a pull request to contribute to this project by following this guideline.

License

This poject is available as open source under the terms of the MIT License.

Originally, this project was based on the visualkeras (under the MIT license), with additional modifications inspired by pytorchviz, and pytorch-summary, both of which are also licensed under the MIT license.

Citation

Please cite this project in your publications if it helps your research.

Note: the paper below describes the API as of its publication date (2024). VisualTorch has since had breaking API changes (see the documentation for the current API) - the DOI always resolves to what was actually reviewed and published, so it isn't updated to match.

@article{Hendria2024,
  doi = {10.21105/joss.06678},
  url = {https://doi.org/10.21105/joss.06678},
  year = {2024},
  publisher = {The Open Journal},
  volume = {9},
  number = {102},
  pages = {6678},
  author = {Willy Fitra Hendria and Paul Gavrikov},
  title = {VisualTorch: Streamlining Visualization for PyTorch Neural Network Architectures},
  journal = {Journal of Open Source Software}
}