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
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.