【PyTorch学习笔记】20：使用nn.Module类及其周边

关于nn.Module

在PyTorch中nn.Module类是用于定义网络中前向结构的父类，当要定义自己的网络结构时就要继承这个类。现有的那些类式接口（如nn.Linear、nn.BatchNorm2d、nn.Conv2d等）也是继承这个类的，nn.Module类可以嵌套若干nn.Module的对象，来形成网络结构的嵌套组合。以下说明使用这个类的一些好处和基本方式。

1 嵌套现有的类

在torch.nn下提供了大量继承了nn.Module类的模块，只要直接使用其初始化函数创建对象，然后调用forward函数就能使用里面的前向计算过程。

2 容器

使用容器模块可以组合多个nn.Module类的对象，最常见的是nn.Sequential容器，它以若若干继承了nn.Module类的对象为参数进行构造，将它们按传入顺序组合起来，在forward时会按这个顺序依次调用，当然要保证前一模块的输出能适合于下一模块的输入才行。

3 模块内部参数管理

使用nn.Module类创建的模块，可以用.parameters()或者.named_parameters()返回 其内的所有参数的迭代器：

from torch import nn

net = nn.Sequential(
    nn.Linear(4, 2),  # 输入维度4,输出维度2的线性层
    nn.Linear(2, 2)  # 输入维度2,输出维度2的线性层
)

print(list(net.parameters()))
print('-' * 10)
print(dict(net.named_parameters()))

运行结果：

[Parameter containing:
tensor([[-0.0632,  0.0654, -0.0377,  0.4113],
        [ 0.2964,  0.4448,  0.3035, -0.0701]], requires_grad=True),
Parameter containing:
tensor([-0.3755, -0.3913], requires_grad=True),
Parameter containing:
tensor([[ 0.0176,  0.6857],
        [-0.1002, -0.2019]], requires_grad=True),
Parameter containing:
tensor([ 0.3695, -0.2610], requires_grad=True)]
----------
{'0.weight': Parameter containing:
tensor([[-0.0632,  0.0654, -0.0377,  0.4113],
        [ 0.2964,  0.4448,  0.3035, -0.0701]], requires_grad=True),
'0.bias': Parameter containing:
tensor([-0.3755, -0.3913], requires_grad=True),
'1.weight': Parameter containing:
tensor([[ 0.0176,  0.6857],
        [-0.1002, -0.2019]], requires_grad=True),
'1.bias': Parameter containing:
tensor([ 0.3695, -0.2610], requires_grad=True)}

其中.named_parameters()还能看到参数名，默认情况下会使用所在的层数+参数类型的方式，从0层开始编号。

另外，使用优化器时，可以用这样的方法，直接将nn.Module类定义的网络参数传进去，而不必手动管理：

optimizer = optim.SGD(net.parameters(), lr=learning_rate)

4 模块树形结构

模块之间通过嵌套组合会形成树形结构，使用.childern()可以获取其直接孩子结点，使用.modules()可以获取其所有子结点。

from torch import nn


class BaseNet(nn.Module):
    """直接使用输入4维输出3维的线性层"""

    def __init__(self):
        super(BaseNet, self).__init__()
        self.net = nn.Linear(4, 3)

    def forward(self, x):
        return self.net(x)


class MyNet(nn.Module):
    """使用Seq容器组合了三个模块"""

    def __init__(self):
        super(MyNet, self).__init__()
        self.net = nn.Sequential(
            BaseNet(),
            nn.ReLU(),
            nn.Linear(3, 2)
        )

    def forward(self, x):
        return self.net(x)


my_net = MyNet()

print(list(my_net.children()))  # 直接孩子
print('-' * 10)
print(list(my_net.modules()))  # 所有孩子

运行结果：

[Sequential(
  (0): BaseNet(
    (net): Linear(in_features=4, out_features=3, bias=True)
  )
  (1): ReLU()
  (2): Linear(in_features=3, out_features=2, bias=True)
)]
----------
[MyNet(
  (net): Sequential(
    (0): BaseNet(
      (net): Linear(in_features=4, out_features=3, bias=True)
    )
    (1): ReLU()
    (2): Linear(in_features=3, out_features=2, bias=True)
  )
), Sequential(
  (0): BaseNet(
    (net): Linear(in_features=4, out_features=3, bias=True)
  )
  (1): ReLU()
  (2): Linear(in_features=3, out_features=2, bias=True)
), BaseNet(
  (net): Linear(in_features=4, out_features=3, bias=True)
), Linear(in_features=4, out_features=3, bias=True), ReLU(), Linear(in_features=3, out_features=2, bias=True)]

可以看到在使用.modules()获取的所有孩子是包括自己这个结点的，然后所有的结点会被放在一个列表里返回。而.childern()只返回自己的直系孩子列表，在这里也就是一个nn.Sequential容器。

5 设备切换

使用.to(device)可以在具体的CPU/GPU上切换，这会将其所有子模块也一起转移过去运行。

注意，模块的.to(device)是原地操作并返回自己的引用，而Tensor的.to(device)不会在当前Tensor上操作，返回的才是在目标设备上对应创建的Tensor。

6 载入和保存检查点

在训练前可以检查是否有检查点文件，使用torch.load()载入检查点文件，然后传入net.load_state_dict()中网络模型设置参数：

net.load_state_dict(torch.load('ckpt.mdl'))

在训练过程中，每隔一定的迭代次数可以保存一下检查点，将当前网络模型的状态传进去：

torch.save(net.state_dict(), 'ckpt.mdl')

7 训练和测试模式的切换

前面学了，如Dropout和Batch Normalization这样的操作在训练和测试中的行为是不同的，为每个模块单独设置训练和测试状态很麻烦，可以直接为网络使用.train()方法切换到训练模式，使用.eval()方法切换到测试模式。

8 实现自定义的层

除了组合各个层以实现自己的网络之外，也可以继承nn.Module类来实现自定义的层。例如一个将数据只保留第一维，其它维合并的Flatten层，可以用在卷积层和全连接层之间：

class Flatten(nn.Module):
    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, input):
        return input.view(input.size(0), -1)

又如DL笔记2里面实现的残差块，也属于自定义的层。

---

如果在继承nn.Module类来实现模块时，出现需要操作Tensor的部分，那么应当使用nn.Parameters将其包装起来，而不应直接使用Tensor。如果直接使用Tensor，那么没法用.parameters()获取到所有参数，也就没法直接传给优化器去记录要优化的包括这些参数了。

class MyLinear(nn.Module):
    def __init__(self, inp, outp):
        super(MyLinear, self).__init__()
        # 线性层的参数w和b,对w而言输出维度放在前面
        self.w = nn.Parameter(torch.randn(outp, inp))
        self.b = nn.Parameter(torch.randn(outp))

    def forward(self, x):
        x = x @ self.w.t() + self.b
        return x

使用nn.Parameters包装Tensor时，自动设置了requires_grad=True，也即默认情况下认为它是反向传播优化的参数之一。