在使用深度学习模型处理图像数据的时候，输入数据的大小在整个网络中是怎么变化的非常重要，但是如果只看代码的话，我们算起来比较麻烦，比如我们经过了各种上采样、下采样等，中间过程可能有几十个网络层，算过来算过去一来是麻烦，二来是不清晰明了。
今天介绍一个用来查看模型概览信息的包，对新手学习非常有帮助！ 这就是torchsummary包（https://github.com/sksq96/pytorch-summary）。
直接上用法。
安装
安装没什么难的，直接用pip即可。
pip install torchsummary
搭建模型
根据自己的需求进行模型的搭建。
模型总体信息概览
比如我们搭建了下面一个网络：
class MyNet(nn.Module):
def __init__(self, inchannels):
super(SRCNN, self).__init__()
self.main = nn.Sequential(
nn.Conv2d(inchannels, 64, kernel_size=9, stride=(1, 1), padding=(4, 4)),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 32, kernel_size=1, stride=(1, 1), padding=(0, 0)),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 1, kernel_size=5, stride=(1, 1), padding=(2, 2))
)
def forward(self, x):
y = self.main(x)
return y
现在我们想看一下输入数据在网络中的大小变化，以及可学习参数的总数量等信息，那么只需要使用summary函数即可：
model = MyNet(inchannels=1).to('cuda')
summary(model, input_size=(1, 222, 222))
然后便可以得到如下信息：
MyNet
从中可以看到整个模型的如下信息：
- 所有的layer
- 每个layer的输出尺寸（-1代表是可变的，这里是batchsize的大小）
- 每一层的参数量
- 总的参数量
- 总得可学习参数量和无需学习的参数量
- 占用空间的大小
有什么用？
目前，个人觉得有很多：
- 最直接的就是可以清楚地了解自己的网络结构；
- 在用卷积神经网络的时候可以用这个来调整kernel size和padding的尺寸，这样子就不用用公式算了；
- 可以来测试自己的模型是不是能跑通。。如果错了的话，上面的结果是出不来的。
参考
【1】https://openbase.com/python/torch-summary
【2】https://clay-atlas.com/us/blog/2020/05/13/pytorch-en-note-torchsummary/
本文由mdnice多平台发布
