PyTorch入门小笔记——利用简单例子观察前向传播各个层输出的size

博主正在学习《深度学习框架PyTorch：入门与实践》，记录一个简单的例子，加深对torch前向传播参数的理解。

这是书籍第二章的一个定义网络的例子，直接看代码可能会不太直观，特别是x = x.view(x.size()[0], -1)这一句，初学者希望能直观感受size的变化，以及fc1中的$16\ast 5\ast 5$的来源：

import torch.nn as nn
import torch.nn.functional as F # 激活和池化都在这里

class Net(nn.Module):
    def __init__(self):
        # nn.Module子类的函数必须在构造函数中执行父类的构造函数
        super().__init__()
        
        self.conv1 = nn.Conv2d(1, 6, 5) # 输入通道1，输出通道6，卷积核5*5，下同
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(x.size()[0], -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
    
net = Net()
print(net)

博主用jupyter记录了每一层输出的size：

# 输入图片的尺寸是32*32，channel为1。即batch*channel*height*width
input = Variable(t.randn(1, 1, 32, 32))
out = net(input)

# 逐层运行并输出每层的size
out1 = net.conv1(input)
pool1 = F.max_pool2d(out1, 2)
out2 = net.conv2(pool1)
pool2 = F.max_pool2d(out2, 2)
stretch = pool2.view(pool2.size()[0], -1)
out3 = net.fc1(strech)
out4 = net.fc2(out3)
out5 = net.fc3(out4)
print('output of conv1:', out1.shape)
print('output of pool1:', pool1.shape)
print('output of conv2:', out2.shape)
print('output of pool2:', pool2.shape)
print('output of stretch:', stretch.shape)
print('output of fc1:', out3.shape)
print('output of fc2:', out4.shape)
print('output of fc3:', out5.shape)

运行结果：

output of conv1: torch.Size([1, 6, 28, 28])
output of pool1: torch.Size([1, 6, 14, 14])
output of conv2: torch.Size([1, 16, 10, 10])
output of pool2: torch.Size([1, 16, 5, 5])
output of strech: torch.Size([1, 400])
output of fc1: torch.Size([1, 120])
output of fc2: torch.Size([1, 84])
output of fc3: torch.Size([1, 10])
​

这样就非常直观了！其中x = x.view(x.size()[0], -1)是将[1, 16, 5, 5]变成了[1, 400]