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
Example usage of make_dot:
model = nn.Sequential()
model.add_module('W0', nn.Linear(8, 16))
model.add_module('tanh', nn.Tanh())
model.add_module('W1', nn.Linear(16, 1))
x = torch.randn(1, 8)
y = model(x)
make_dot(y.mean(), params=dict(model.named_parameters()))
Set show_attrs=True and show_saved=True to see what autograd saves for the backward pass. (Note that this is only available for pytorch >= 1.9.)
model = nn.Sequential()
model.add_module('W0', nn.Linear(8, 16))
model.add_module('tanh', nn.Tanh())
model.add_module('W1', nn.Linear(16, 1))
x = torch.randn(1, 8)
y = model(x)
make_dot(y.mean(), params=dict(model.named_parameters()), show_attrs=True, show_saved=True)
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. Other contributors are @willprice, @soulitzer, @albanD.