How do I visualize a net in Pytorch?

Question

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision.models as models
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torch.autograd import Variable
from torchvision.models.vgg import model_urls
from torchviz import make_dot

batch_size = 3
learning_rate =0.0002
epoch = 50

resnet = models.resnet50(pretrained=True)
print resnet
make_dot(resnet)

I want to visualize resnet from the pytorch models. How can I do it? I tried to use torchviz but it gives an error:

'ResNet' object has no attribute 'grad_fn'

Answer 1

You can have a look at PyTorchViz (https://github.com/szagoruyko/pytorchviz), "A small package to create visualizations of PyTorch execution graphs and traces."

Answer 2

Here is how you do it with torchviz if you want to save the image:

# http://www.bnikolic.co.uk/blog/pytorch-detach.html

import torch
from torchviz import make_dot

x=torch.ones(10, requires_grad=True)
weights = {'x':x}

y=x**2
z=x**3
r=(y+z).sum()

make_dot(r).render("attached", format="png")

screenshot of image you get:

source: http://www.bnikolic.co.uk/blog/pytorch-detach.html

Answer 3

You can use TensorBoard for visualization. TensorBoard is now fully supported in PyTorch version 1.2.0. More info: https://pytorch.org/docs/stable/tensorboard.html

Answer 4

make_dot expects a variable (i.e., tensor with grad_fn), not the model itself.
try:

x = torch.zeros(1, 3, 224, 224, dtype=torch.float, requires_grad=False)
out = resnet(x)
make_dot(out)  # plot graph of variable, not of a nn.Module