/pytorchviz

A small package to create visualizations of PyTorch execution graphs

Primary LanguageJupyter NotebookMIT LicenseMIT

PyTorchViz

A small package to create visualizations of PyTorch execution graphs and traces.

Installation

Install graphviz, e.g.:

brew install graphviz

Install the package itself:

pip install torchviz

Usage

screen shot 2018-01-30 at 16 13 01

There are two functions, make_dot to make graphs from any PyTorch functions (requires that at least one input Variable requires_grad), and make_dot_from_trace that uses outputs of torch.jit.trace (does not always work). See examples.ipynb.

Acknowledgements

The script was moved from functional-zoo where it was created with the help of Adam Paszke, Soumith Chintala, Anton Osokin, and uses bits from tensorboard-pytorch.