Tensorboard: How to view pytorch model summary?

Question

I have the following network.

import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter

class Net(nn.Module):
    def __init__(self,input_shape, num_classes):
        super(Net, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=(4,4)),

            nn.Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=(4,4)),
        )

        x = self.conv(torch.rand(input_shape))
        in_features = np.prod(x.shape)

        self.classifier = nn.Sequential(
            nn.Linear(in_features=in_features, out_features=num_classes),
        )

    def forward(self, x):
        x = self.feature_extractor(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

net = Net(input_shape=(1,64,1292), num_classes=4)
print(net)

This prints the following:-

Net(
  (conv): Sequential(
    (0): Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): MaxPool2d(kernel_size=(4, 4), stride=(4, 4), padding=0, dilation=1, ceil_mode=False)
    (3): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (4): ReLU(inplace=True)
    (5): MaxPool2d(kernel_size=(4, 4), stride=(4, 4), padding=0, dilation=1, ceil_mode=False)
  )
  (classifier): Sequential(
    (0): Linear(in_features=320, out_features=4, bias=True)
  )
)

However, I am trying various experiments and I want to keep track of network architecture on Tensorboard. I know there is a function writer.add_graph(model, input_to_model) but it requires input, or at least its shape should be known.

So, I tried writer.add_text("model", str(model)), but formatting is screwed up in tensorboard.

My question is, is there a way to at least visualize the way I can see by using print function in the tensorboard?

Answer 1

I can see everything is going right but there is just a formatting issue. Tensorboard understands markdown so you can actually replace \n with <br/> and with  .

Here is a detailed walkthrough. Suppose you have the following model:-

import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter

class Net(nn.Module):
    def __init__(self,input_shape, num_classes):
        super(Net, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=(4,4)),

            nn.Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=(4,4)),
        )

        x = self.conv(torch.rand(input_shape))
        in_features = np.prod(x.shape)

        self.classifier = nn.Sequential(
            nn.Linear(in_features=in_features, out_features=num_classes),
        )

    def forward(self, x):
        x = self.feature_extractor(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

net = Net(input_shape=(1,64,1292), num_classes=4)
print(net)

This prints the following and if can actually show it in the Tensorboard.

Net(
  (conv): Sequential(
    (0): Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): MaxPool2d(kernel_size=(4, 4), stride=(4, 4), padding=0, dilation=1, ceil_mode=False)
    (3): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (4): ReLU(inplace=True)
    (5): MaxPool2d(kernel_size=(4, 4), stride=(4, 4), padding=0, dilation=1, ceil_mode=False)
  )
  (classifier): Sequential(
    (0): Linear(in_features=320, out_features=4, bias=True)
  )
)

There is function in add_graph(model, input) in SummaryWriter but you must create dummy input and in some cases it is difficult of to always know them. Instead do following:-

writer = SummaryWriter()

model_summary = str(model).replace( '\n', '<br/>').replace(' ', '&nbsp;')
writer.add_text("model", model_summary)

writer.close()

Above produces following text in tensorboard:-