GIS与Climate

V1

2022/07/30阅读:27主题:默认主题

深度学习模型训练前的必做工作:总览模型信息

在使用深度学习模型处理图像数据的时候,输入数据的大小在整个网络中是怎么变化的非常重要,但是如果只看代码的话,我们算起来比较麻烦,比如我们经过了各种上采样、下采样等,中间过程可能有几十个网络层,算过来算过去一来是麻烦,二来是不清晰明了。

今天介绍一个用来查看模型概览信息的包,对新手学习非常有帮助! 这就是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
MyNet

从中可以看到整个模型的如下信息:

  • 所有的layer
  • 每个layer的输出尺寸(-1代表是可变的,这里是batchsize的大小)
  • 每一层的参数量
  • 总的参数量
  • 总得可学习参数量和无需学习的参数量
  • 占用空间的大小

有什么用?

目前,个人觉得有很多:

  1. 最直接的就是可以清楚地了解自己的网络结构;
  2. 在用卷积神经网络的时候可以用这个来调整kernel size和padding的尺寸,这样子就不用用公式算了;
  3. 可以来测试自己的模型是不是能跑通。。如果错了的话,上面的结果是出不来的。

参考

【1】https://openbase.com/python/torch-summary
【2】https://clay-atlas.com/us/blog/2020/05/13/pytorch-en-note-torchsummary/

分类:

人工智能

标签:

深度学习

作者介绍

GIS与Climate
V1

公众号:GIS与Climate,欢迎关注