PyTorch学习笔记-神经网络模型搭建小实战

1. torch.nn.Sequential

torch.nn.Sequential 是一个 Sequential 容器，能够在容器中嵌套各种实现神经网络中具体功能相关的类，来完成对神经网络模型的搭建。模块的加入一般有两种方式，一种是直接嵌套，另一种是以 OrderedDict 有序字典的方式进行传入，这两种方式的唯一区别是：

使用 OrderedDict 搭建的模型的每个模块都有我们自定义的名字。
直接嵌套默认使用从零开始的数字序列作为每个模块的名字。

（1）直接嵌套方法的代码如下：

import torch.nn as nn

model = nn.Sequential(
    nn.Conv2d(1, 20, 5),
    nn.ReLU(),
    nn.Conv2d(20, 64, 5),
    nn.ReLU()
)

print(model)
# Sequential(
#   (0): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
#   (1): ReLU()
#   (2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
#   (3): ReLU()
# )

（2）使用 OrderedDict 的代码如下：

import torch.nn as nn
from collections import OrderedDict

model = nn.Sequential(OrderedDict([
    ('Conv1', nn.Conv2d(1, 20, 5)),
    ('ReLU1', nn.ReLU()),
    ('Conv2', nn.Conv2d(20, 64, 5)),
    ('ReLU2', nn.ReLU())
]))

print(model)
# Sequential(
#   (Conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
#   (ReLU1): ReLU()
#   (Conv2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
#   (ReLU2): ReLU()
# )

2. 小实战

现在我们来搭建下图所示的很简单的神经网络模型：

在这里插入图片描述

对模型进行分析可以看到第一层为卷积层 conv1，其卷积核 kernel_size=5，处理后的图像通道数从3变成了32，说明 out_channels=32，且图像大小没有变化，通过计算可知 stride=1, padding=2；第二层为最大池化层，其池化核 kernel_size=2，处理后的图像大小变为原来的一半。同理可以分析出之后每一层的参数。

代码如下：

from torch.utils.tensorboard import SummaryWriter
import torch.nn as nn
import torch

class CIFAR10_Network(nn.Module):
    def __init__(self):
        super(CIFAR10_Network, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),  # [32, 32, 32]
            nn.MaxPool2d(kernel_size=2),  # [32, 16, 16]
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),  # [32, 16, 16]
            nn.MaxPool2d(kernel_size=2),  # [32, 8, 8]
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),  # [64, 8, 8]
            nn.MaxPool2d(kernel_size=2),  # [64, 4, 4]
            nn.Flatten(),  # [1024]
            nn.Linear(in_features=1024, out_features=64),  # [64]
            nn.Linear(in_features=64, out_features=10) # [10]
        )

    def forward(self, input):
        output = self.model(input)
        return output

network = CIFAR10_Network()

input = torch.randn(64, 3, 32, 32)  # 返回一个包含了从标准正态分布中抽取的一组随机数的张量
print(input.shape)  # torch.Size([64, 3, 32, 32])
output = network(input)
print(output.shape)  # torch.Size([64, 10])

writer = SummaryWriter('logs')
writer.add_graph(network, input)  # 生成计算图
writer.close()

使用 add_graph 函数可以在 TensorBoard 中生成神经网络的计算图，通过计算图可以很清晰地看到每一层计算时数据流入流出的结果，打开 TensorBoard 看一下结果：

在这里插入图片描述

双击相应的标签可以进一步深入查看更详细的信息，例如我们将我们构建的这个神经网络的类展开可以看到其中的模型：

在这里插入图片描述

将这个模型继续展开后就能看到我们在 Sequential 中定义的网络的各个层：

在这里插入图片描述

最后展开某一层即可看到在这层的数据流动情况：

在这里插入图片描述