A small package to create visualizations of PyTorch execution graphs and traces.
Install graphviz, e.g.:
brew install graphviz Install the package itself:
pip install torchviz 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) 
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.