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

当设计神经网络时，常常被每层的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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

此外，如果你的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)
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

最终输出是一样的（输出有点长，这里就放部分的截图）：

这样就可以根据输出的shape来设计和添加合适的层了。

参考：

• LazyModuleMixin — PyTorch 1.12 documentation
• 7.1. 深度卷积神经网络（AlexNet） — 动手学深度学习 2.0.0-beta1 documentation