不知道你有没有查看模型各层输出尺寸的需要，今天介绍一个简单的库来达到这个目的。

##torchsummary
torchsummary这个工具可以用于模型的可视化，会输出模型各层的详细参数以及各层输出的尺寸。有的时候我们还是很在乎每层输出的shape的，用这个就比较容易查看了，用起来很方便。

##torchsummary的安装
pip install torchsummary

##torchsummary的使用

####torchsummary使用起来很简单，直接调用summary()，传入你的model以及model输入tensor的尺寸就可以了，因为在这里是做了一个前向的运算的

1
2
from torchsummary import summary
summary(your_model, input_size=(channels, H, W))

##小例子

##torchvision.models的model

1
2
3
4
5
6
7
8
import torch
from torchsummary import summary
from torchvision.models import vgg11

model = vgg11(pretrained=False)
if torch.cuda.is_available():
    model.cuda()
summary(model, (3, 224, 224))

输出
这里的output shape就是每层输出的tensor的shape

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 224, 224]            1792
              ReLU-2         [-1, 64, 224, 224]               0
         MaxPool2d-3         [-1, 64, 112, 112]               0
            Conv2d-4        [-1, 128, 112, 112]           73856
              ReLU-5        [-1, 128, 112, 112]               0
         MaxPool2d-6          [-1, 128, 56, 56]               0
            Conv2d-7          [-1, 256, 56, 56]          295168
              ReLU-8          [-1, 256, 56, 56]               0
            Conv2d-9          [-1, 256, 56, 56]          590080
             ReLU-10          [-1, 256, 56, 56]               0
        MaxPool2d-11          [-1, 256, 28, 28]               0
           Conv2d-12          [-1, 512, 28, 28]         1180160
             ReLU-13          [-1, 512, 28, 28]               0
           Conv2d-14          [-1, 512, 28, 28]         2359808
             ReLU-15          [-1, 512, 28, 28]               0
        MaxPool2d-16          [-1, 512, 14, 14]               0
           Conv2d-17          [-1, 512, 14, 14]         2359808
             ReLU-18          [-1, 512, 14, 14]               0
           Conv2d-19          [-1, 512, 14, 14]         2359808
             ReLU-20          [-1, 512, 14, 14]               0
        MaxPool2d-21            [-1, 512, 7, 7]               0
           Linear-22                 [-1, 4096]       102764544
             ReLU-23                 [-1, 4096]               0
          Dropout-24                 [-1, 4096]               0
           Linear-25                 [-1, 4096]        16781312
             ReLU-26                 [-1, 4096]               0
          Dropout-27                 [-1, 4096]               0
           Linear-28                 [-1, 1000]         4097000
================================================================
Total params: 132863336
Trainable params: 132863336
Non-trainable params: 0
----------------------------------------------------------------

###小例子2

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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)

model = Net()
if torch.cuda.is_available():
    model.cuda()

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

输出

1
2
3
4
5
6
7
8
9
10
11
12
13
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 10, 24, 24]             260
            Conv2d-2             [-1, 20, 8, 8]            5020
         Dropout2d-3             [-1, 20, 8, 8]               0
            Linear-4                   [-1, 50]           16050
            Linear-5                   [-1, 10]             510
================================================================
Total params: 21840
Trainable params: 21840
Non-trainable params: 0
----------------------------------------------------------------

###地址在这里
https://github.com/sksq96/pytorch-summary