Network related parameters in Pytorch show [feature map] and model flops calculation and reasoning time calculation

1. Network related parameter display

(1) Visualization of the network structure, as shown in the red box below

(2) Feature map visualization and parameter calculation

Need to install the torchsummary package:

pip install torchsummary

Examples are as follows:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # PyTorch v0.4.0
model = Net().to(device)

summary(model, (1, 28, 28))

The results are as follows: 

#         每一层类型               特征图shape         每层的参数量
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 10, 24, 24]             260
            Conv2d-2             [-1, 20, 8, 8]           5,020
         Dropout2d-3             [-1, 20, 8, 8]               0
            Linear-4                   [-1, 50]          16,050
            Linear-5                   [-1, 10]             510
================================================================
Total params: 21,840       # 模型整体的参数量（上面层参数量相加）
Trainable params: 21,840 
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00     # 图片预处理后的大小
Forward/backward pass size (MB): 0.06   # 正向/反向传播一次的内存大小
Params size (MB): 0.08
Estimated Total Size (MB): 0.15
--------------------------------------------------------------

Reference link:

https://github.com/sksq96/pytorch-summary

Second, the flops calculation of the model

(1) Install thop

pip install thop

(2) Basic use

from torchvision.models import resnet50
from thop import profile
model = resnet50()
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ))

Reference: https://www.jianshu.com/p/6514b8fb1ada

          https://zhuanlan.zhihu.com/p/337810633

Focus on reference: https://blog.csdn.net/junmuzi/article/details/83109660?utm_medium=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-2.control&depth_1-utm_source=distribute.pc_relevant.none-task-blog -BlogCommendFromMachineLearnPai2-2.control

3. The average inference time of the model

   from torchvision.models import googlenet
   from thop import profile
   model = googlenet()
   input = torch.randn(1, 3, 224, 224)
   flops, params = profile(model, inputs=(input,))
   print("flops：",flops)
   print('params',params)

Reference: https://www.jianshu.com/p/cbada26ea29d?from=groupmessage