torchsummary：计算神经网络模型各层输出特征图尺寸及参数量

之前写过一篇自动计算模型参数量、FLOPs、乘加数以及所需内存等数据的博客，介绍了torchstat的用法。现介绍一款更为轻量的工具：torchsummary。使用方法如下：

1：安装

pip install torchsummary

2：导入和使用

【注意】：此工具是针对PyTorch的，需配合PyTorch使用！
使用顺序可概括如下：
（1）导入torchsummary中的summary对象；
（2）建立神经网络模型；
（3）输入 模型（model）、输入尺寸（input_size）、批次大小（batch_size）、运行平台（device）信息，运行后即可得到summary函数的返回值。

summary函数的接口信息如下：

summary(model, input_size, batch_size, device)

4个参数在（3）中已进行了解释，其中device是指cpu或gpu.

3：使用实例

import torch
import torchvision
# 导入torchsummary
from torchsummary import summary

# 需要使用device来指定网络在GPU还是CPU运行
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 建立神经网络模型，这里直接导入已有模型
# model = model().to(device)
model = torchvision.models.vgg11_bn().to(device)
# 使用summary，注意输入维度的顺序
summary(model, input_size=(3, 224, 224))

输出如下：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 224, 224]           1,792
       BatchNorm2d-2         [-1, 64, 224, 224]             128
              ReLU-3         [-1, 64, 224, 224]               0
         MaxPool2d-4         [-1, 64, 112, 112]               0
            Conv2d-5        [-1, 128, 112, 112]          73,856
       BatchNorm2d-6        [-1, 128, 112, 112]             256
              ReLU-7        [-1, 128, 112, 112]               0
         MaxPool2d-8          [-1, 128, 56, 56]               0
            Conv2d-9          [-1, 256, 56, 56]         295,168
      BatchNorm2d-10          [-1, 256, 56, 56]             512
             ReLU-11          [-1, 256, 56, 56]               0
           Conv2d-12          [-1, 256, 56, 56]         590,080
      BatchNorm2d-13          [-1, 256, 56, 56]             512
             ReLU-14          [-1, 256, 56, 56]               0
        MaxPool2d-15          [-1, 256, 28, 28]               0
           Conv2d-16          [-1, 512, 28, 28]       1,180,160
      BatchNorm2d-17          [-1, 512, 28, 28]           1,024
             ReLU-18          [-1, 512, 28, 28]               0
           Conv2d-19          [-1, 512, 28, 28]       2,359,808
      BatchNorm2d-20          [-1, 512, 28, 28]           1,024
             ReLU-21          [-1, 512, 28, 28]               0
        MaxPool2d-22          [-1, 512, 14, 14]               0
           Conv2d-23          [-1, 512, 14, 14]       2,359,808
      BatchNorm2d-24          [-1, 512, 14, 14]           1,024
             ReLU-25          [-1, 512, 14, 14]               0
           Conv2d-26          [-1, 512, 14, 14]       2,359,808
      BatchNorm2d-27          [-1, 512, 14, 14]           1,024
             ReLU-28          [-1, 512, 14, 14]               0
        MaxPool2d-29            [-1, 512, 7, 7]               0
AdaptiveAvgPool2d-30            [-1, 512, 7, 7]               0
           Linear-31                 [-1, 4096]     102,764,544
             ReLU-32                 [-1, 4096]               0
          Dropout-33                 [-1, 4096]               0
           Linear-34                 [-1, 4096]      16,781,312
             ReLU-35                 [-1, 4096]               0
          Dropout-36                 [-1, 4096]               0
           Linear-37                 [-1, 1000]       4,097,000
================================================================
Total params: 132,868,840
Trainable params: 132,868,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 182.03
Params size (MB): 506.85
Estimated Total Size (MB): 689.46
----------------------------------------------------------------

可以看出，batch_size可以不指定，默认为-1。summary函数会对模型中的每层输出特征图尺寸进行计算，并计算每层含有的参数量以及模型的参数总量等信息，对于逐层统计计算和分析非常直观和简洁。