[PYTHON] About the matter that torch summary can be really used when building a model with Pytorch

Introduction

I'm building my own model, and before I connect it to the fully connected layer, I always ask, "What are the features of the input?" You can understand the structure of the model by typing print (model) often, but you cannot check the size of the Feature Map. That's where the torch summary comes in handy.

Who is the torch summary?

Simply put, it's a you can see the size of the feature map.

How to use torch summary

First, make a model

This time I made the following simple model. I haven't written it until I classify it.

Convolution ➡︎BN➡︎ReLU➡︎pooling➡︎ Convolution ➡︎BN➡︎ReLU➡︎pooling➡︎ Convolution ➡︎ Global Average Pooling

                                                                                                                                     
import torch
import torch.nn as nn
   
class SimpleCNN(nn.Module):     
    def __init__(self):
        super(SimpleCNN,self).__init__()
 
        self.conv1 = nn.Conv2d(3,16,kernel_size=3,stride=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d((2,2))
        self.conv2 = nn.Conv2d(16,32,kernel_size=3,stride=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.conv3 = nn.Conv2d(32,64,kernel_size=3,stride=1)
        self.gap = nn.AdaptiveMaxPool2d(1)
  
    def forward(self,x):  
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x) 
        x = self.maxpool(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.conv3(x)
        x = self.gap(x)
   
        return x

Installation of torch summary

pip install torchsummary

torchsummary How to use

from torchsummary import summary
model = SimpleCNN()  
summary(model,(3,224,224)) # summary(model,(channels,H,W))

This time, I am trying to assume an image input size of 224x224. If you want to try other resolutions, change the values of H and W.

Summary output

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 16, 222, 222]             448
       BatchNorm2d-2         [-1, 16, 222, 222]              32
              ReLU-3         [-1, 16, 222, 222]               0
         MaxPool2d-4         [-1, 16, 111, 111]               0
            Conv2d-5         [-1, 32, 109, 109]           4,640
       BatchNorm2d-6         [-1, 32, 109, 109]              64
              ReLU-7         [-1, 32, 109, 109]               0
         MaxPool2d-8           [-1, 32, 54, 54]               0
            Conv2d-9           [-1, 64, 52, 52]          18,496
AdaptiveMaxPool2d-10             [-1, 64, 1, 1]               0
================================================================
Total params: 23,680
Trainable params: 23,680
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 30.29
Params size (MB): 0.09
Estimated Total Size (MB): 30.95
----------------------------------------------------------------

It is quite convenient to be able to check the Shape of Output. I'm grateful that it also counts the number of parameters.

At the end

torchsummary is convenient, so please use it.