A small package to create visualizations of PyTorch execution graphs and traces.
pip install git+https://github.com/szagoruyko/pytorchviz
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.
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.