Pytorch神经网络设计技巧-如何得到各层的input shape

wonder-wall

已于 2022-09-23 19:14:48 修改

阅读量1.5k

点赞数 3

分类专栏：小技巧 Pytorch

于 2022-08-31 11:57:09 首次发布

本文链接：https://blog.csdn.net/qq_41129489/article/details/126620861

版权

小技巧同时被 2 个专栏收录

8 篇文章 6 订阅

订阅专栏

Pytorch

1 篇文章 0 订阅

订阅专栏

当设计神经网络时，常常被每层的input shape所烦恼，计算也很恼人。
在使用pytorch设计的时候，有两种方法能够协助你进行输入维度的设置。

方法一：懒加载（nn.Lazy**）

Pytorch 提供了各个常用层的Lazy版本，类似于懒加载的思想
这些层有懒初始化（lazy initialization）的机制，在使用时不用显示指定input shape
它们会在第一次前向传播的时候获取输入数据的shape来初始化自己的input shape
省去人工计算input shape的操作。
在这里插入图片描述

方法二：利用前向传播调试

懒初始化是Pytorch比较新的特性，官方也提示后续会对懒初始化机制进行不断地完善。
如果坚持使用不带lazy的层，可以在设计的时候给一个输入数据
让数据在网络中进行一次前向传播，打印每个层输出的shape：

# 生成1幅单通道的224*224大小的随机图像
X = torch.randn(1, 1, 224, 224)
# 这里的net是nn.Sequential类型，里面包含了各个layer，它是可迭代的结构。
net = nn.Sequential(
    nn.Conv2d(in_channels=1, out_channels=96, kernel_size=11, stride=4, padding=1), nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, padding=2), nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, padding=1), nn.ReLU(),
    nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, padding=1), nn.ReLU(),
    nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, padding=1), nn.ReLU(),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Flatten(),
    nn.Linear(in_features=6400, out_features=4096) , nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(in_features=4096, out_features=4096) , nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(in_features=4096, out_features=10)
)
for layer in net:
    X=layer(X)
    print(layer.__class__.__name__,'output shape:\t',X.shape)

此外，如果你的net是继承自nn.Module的类实例，可以设置一个类型为nn.Sequential的成员变量封装各个layer，然后for循环来遍历这个成员变量，本质是一样的：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=96, kernel_size=11, stride=4, padding=1), nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, padding=2), nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, padding=1), nn.ReLU(),
            nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, padding=1), nn.ReLU(),
            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, padding=1), nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Flatten(),
            nn.Linear(in_features=6400, out_features=4096), nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=4096, out_features=4096), nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=4096, out_features=10)
        )
    def forward(self, x):
        return self.layers(x)
net = Net()
X = torch.randn(1, 1, 224, 224)
for layer in net.layers:
    X = layer(X)
    print(layer.__class__.__name__, 'output shape:\t', X.shape)