How to generate the architectural figure of the network in pytorch when the type of the output of the network is list?

Question

I wonder how to use tochviz to generate network architecture, when the output is a list type? the demo code is as follows:

 import torch
 import torch.nn as nn
 class ConvNet(nn.Module):
     def __init__(self):
         super(ConvNet, self).__init__()
         self.conv1 = nn.Sequential(
             nn.Conv2d(1, 16, 3, 1, 1),
             nn.ReLU(),
             nn.AvgPool2d(2, 2)
         )
         self.conv2 = nn.Sequential(
             nn.Conv2d(16, 32, 3, 1, 1),
             nn.ReLU(),
             nn.MaxPool2d(2, 2)
         )
         self.fc = nn.Sequential(
             nn.Linear(32 * 7 * 7, 128),
             nn.ReLU(),
             nn.Linear(128, 64),
             nn.ReLU()
         )
         self.out = nn.Linear(64, 10)
     def forward(self, x):
         x = self.conv1(x)
         x = self.conv2(x)
         x = x.view(x.size(0), -1)
         x = self.fc(x)
         output = []
         output.append(x)
         output.append(self.out(x))
         return output
 MyConvNet = ConvNet()

and I use torchviz to view this network's architecture like

 from torchviz import make_dot
 x = torch.randn(1, 1, 28, 28).requires_grad_(True)
 y = MyConvNet(x)   
 MyConvNetVis = make_dot(y, params=dict(list(MyConvNet.named_parameters()) + [('x', x)]))
 MyConvNetVis.format = "png"
 MyConvNetVis.directory = "data"
 MyConvNetVis.view()

then, I was blocked with this problem

AttributeError                            Traceback (most recent call last)
<ipython-input-23-c8e3cd3a8b4e> in <module>
      2 x = torch.randn(1, 1, 28, 28).requires_grad_(True)
      3 y = MyConvNet(x)
----> 4 MyConvNetVis = make_dot(y, params=dict(list(MyConvNet.named_parameters()) + [('x', x)]))
      5 MyConvNetVis.format = "png"
      6 MyConvNetVis.directory = "data"

~/anaconda3/envs/torch1.3/lib/python3.6/site-packages/torchviz/dot.py in make_dot(var, params)
     35         return '(' + (', ').join(['%d' % v for v in size]) + ')'
     36 
---> 37     output_nodes = (var.grad_fn,) if not isinstance(var, tuple) else tuple(v.grad_fn for v in var)
     38 
     39     def add_nodes(var):

AttributeError: 'list' object has no attribute 'grad_fn'

Any advice will be appreciated.

Answer 1

The error indicates torchviz is trying to navigate through the network using grad_fn in order to compute its own graph. However, a tuple is not a tensor and doesn't have the gra_fn property. I'm not quite sure you can have multiple outputs (i.e. a tuple as output) when working with torchviz. As a workaround, if you just want to visualize the network you could replace the tuple with concatenation of the two tensors using torch.cat:

def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)
    x = x.view(x.size(0), -1)
    x = self.fc(x)
    out = self.out(x)
    output = torch.cat([x, out], dim=1)
    return output

The result is:

Notice how the last node is a CatBackward with two incoming branches, one from AddmmBackward (out) and the other from ReluBackward0 (x). This last node is fictitious and is not present in your actual model, so you could remove it from the graph, by hand.

Answer 2

earlier answer shows packages that can build the architectural diagram/graph for a Pytorch Model:

• torchviz/pytorchviz
• TensorBoard
• Netron
• HiddenLayer