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.
Simply put, it's a you can see the size of the feature map
.
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
pip install torchsummary
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
.
----------------------------------------------------------------
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.
torchsummary is convenient, so please use it.
Recommended Posts