PyTorch深度学习实战（10）—— 神经网络工具箱nn.Module
虽然使用autograd可以实现深度学习模型，但是它的抽象程度较低，需要编写的代码量极大。本文介绍的nn模块，是构建于autograd之上的神经网络模块。除了nn，本章还会介绍神经网络中常用的工具，比如优化器optim、初始化init等。
torch.nn是专门为深度学习设计的模块，它的核心数据结构是Module。torch.nn是一个抽象的概念，既可以表示神经网络中的某个层（layer），又可以表示一个包含很多层的神经网络。在实际使用中，最常见的做法是继承nn.Module，然后编写自己的网络/层。下面先来看看如何使用nn.Module实现自己的全连接层。全连接层，又名仿射层，它的输出y和输入x满足y = Wx+b ，其中w和b是可学习的参数。
In: import torch as t
from torch import nn
print(t.__version__)
Out:1.8.0
In: # 继承nn.Module，必须重写构造函数__init__和前向传播函数forward
class Linear(nn.Module):
def __init__(self, in_features, out_features):
super().__init__() # 等价于nn.Module.__init__(self)，常用super方式
# nn.Parameter内的参数是网络中可学习的参数
self.W = nn.Parameter(t.randn(in_features, out_features))
self.b = nn.Parameter(t.randn(out_features))
def forward(self, x):
x = x.mm(self.W) # 矩阵乘法，等价于x@(self.W)
return x + self.b.expand_as(x)
In: layer = Linear(4,3)
input = t.randn(2,4)
output = layer(input)
output
Out:tensor([[-1.0987, -0.2932, -3.5264],
[-0.0662, -5.5573, -8.1498]], grad_fn=<AddBackward0>)
In: for name, parameter in layer.named_parameters():
print(name, parameter) # W和b
Out:
W Parameter containing:
tensor([[ 0.5180, 1.4337, 0.4373],
[ 0.2299, -1.6198, -0.7570],
[ 0.0694, -1.7724, -0.2443],
[ 0.0258, 0.1944, 3.4072]], requires_grad=True)
b Parameter containing:
tensor([-0.4774, 1.4022, -1.4314], requires_grad=True)
从上面的例子可以看出，全连接层的实现非常简单，代码量不超过10行。以上述代码为例，在自定义层时需要注意以下几点。
- 自定义层Linear必须继承nn.Module，在构造函数中需要调用nn.Module的构造函数，即super().__init__() 或nn.Module.__init__(self)，笔者推荐使用第一种方法。
- 在构造函数__init__()中必须自行定义可学习的参数，并封装成nn.Parameter。在本例中将w和b封装成Parameter。Parameter是一种特殊的Tensor，它默认需要求导（requires_grad=True），感兴趣的读者可以通过nn.Parameter??查看Parameter类的源代码。
- forward函数实现了前向传播过程，它的输入可以是一个或者多个Tensor。
- 反向传播函数无需手动编写，nn.Module能够利用autograd自动实现反向传播，这点比Function简单许多。
- 使用时，可以将layer看成数学概念中的函数，调用layer(input)可以得到input对应的结果，它等价于layers.__call__(input)。在__call__函数中，主要是调用 layer.forward(x)，同时还对钩子函数（hook）做了一些处理。在实际使用中应尽量使用layer(x)，而不是使用layer.forward(x)，关于钩子技术的具体内容将在下文讲解。
- nn.Module中的可学习参数可以通过named_parameters()或者parameters()返回一个迭代器，前者会给每个参数都附上名字，使其更具有辨识度。
利用nn.Module实现的全连接层，相较于利用Function实现的更加简单，这是因为无需手动编写反向传播函数。nn.Module能够自动检测到自己的Parameter，并将其作为学习参数。除了Parameter，module还可能包含子module，主module能够递归查找子module中的Parameter，下面以多层感知机为例进行说明。
多层感知机的网络结构如图4-1所示，它由两个全连接层组成，采用sigmoid函数作为激活函数。其中x表示输入，y表示输出，b表示偏置，W表示全连接层的参数。
In: class Perceptron(nn.Module):
def __init__(self, in_features, hidden_features, out_features):
super().__init__()
# 此处的Linear是前面自定义的全连接层
self.layer1 = Linear(in_features, hidden_features)
self.layer2 = Linear(hidden_features, out_features)
def forward(self, x):
x = self.layer1(x)
x = t.sigmoid(x)
return self.layer2(x)
In: perceptron = Perceptron(3, 4, 1)
for name, param in perceptron.named_parameters():
print(name, param.size())
Out:layer1.W torch.Size([3, 4])
layer1.b torch.Size([4])
layer2.W torch.Size([4, 1])
layer2.b torch.Size([1])
即使是稍微复杂的多层感知机，它的实现仍然很简单。在构造函数__init__()中，可以将前面自定义的Linear层（module）作为当前module对象的一个子module。子module的可学习参数，也会成为当前module的可学习参数。在forward函数中，可以加上各层之间的处理函数（如激活函数、数学处理等），并定义层与层之间的关系。
在module中，Parameter的全局命名规范如下。
- Parameter直接命名。例如self.param_name = nn.Parameter(t.randn(3, 4))，可以直接命名为param_name。
- 子module中的Parameter，会在其名字前面加上当前module的名字。例如self.sub_module = SubModel()，在SubModel中有个Parameter的名字叫做param_name，那么二者拼接而成的参数名称是sub_module.param_name。
为了方便用户使用，PyTorch实现了神经网络中绝大多数的网络层，这些层都继承于nn.Module，它们都封装了可学习参数Parameter，并实现了forward函数。同时，大部分layer都专门针对GPU运算进行了cuDNN优化，它们的速度和性能都十分优异。本章不会对nn.Module中的所有层进行详细介绍，具体内容读者可参考官方文档，或在IPython/Jupyter中使用nn.layer?进行查看。读者在阅读文档时应该主要关注以下几点。
- 构造函数的参数，例如nn.Linear(in_features, out_features, bias)，需关注这三个参数的作用。
- 属性、可学习的网络参数和包含的子module。如nn.Linear中有weight和bias两个可学习参数，不包含子module。
- 输入、输出的形状，例如nn.linear的输入形状是(N, input_features)，输出为(N, output_features)，其中N是batch_size。
这些自定义layer对输入形状都有规定：输入的是一个batch数据，而不是单个数据。当输入只有一个数据时，必须调用tensor.unsqueeze(0)或tensor[None]将数据伪装成batch_size=1的一个batch。
下面将从应用层面出发，对一些常用的网络层进行简单介绍。