模型容器

Torch.nn中一个重要的概念是模型容器 (Containers)，常用的容器有 3 个，这些容器都是继承自nn.Module。

  • nn.Sequetial：常用于 block 构建，按照顺序包装多个网络层;
  • nn.ModuleList：常用于大量重复网络构建，通过 for 循环实现重复构建，像 python 的 list 一样包装多个网络层，可以迭代;
  • nn.ModuleDict：常用于可选择的网络层，像 python 的 dict一样包装多个网络层，通过 (key, value) 的方式为每个网络层指定名称;

nn.Sequential

在传统的机器学习中，有一个步骤是特征工程，我们需要从数据中人为地提取特征，然后把特征输入到分类器中预测。在深度学习的时代，特征工程的概念被弱化了，特征提取和分类器这两步被融合到了一个神经网络中。在卷积神经网络中，前面的卷积层以及池化层可以认为是特征提取部分，而后面的全连接层可以认为是分类器部分。比如 LeNet 就可以分为特征提取和分类器两部分，这 2 部分都可以分别使用 nn.Seuqtial 来包装。

image-20220912181620491

代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class LeNetSequetial(nn.Module):
    def __init__(self, classes):
        super(LeNet2, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.AvgPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.AvgPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)
        return x

在初始化时，nn.Sequetial会调用__init__()方法，将每一个子 module 添加到自身_modules属性中。这里可以看到，我们传入的参数可以是一个 list，或者一个 OrderDict。如果是一个 OrderDict，那么则使用 OrderDict 里的 key，否则使用数字作为 key (OrderDict 的情况会在下面提及)。

1
2
3
4
5
6
7
8
def __init__(self, *args):
        super(Sequential, self).__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            for key, module in args[0].items():
                self.add_module(key, module)
        else:
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)

网络初始化完成后有两个子 module：features和classifier。

image-20220912181635902

而features中的子 module 如下，每个网络层以序号作为 key：

image-20220912181654578

在进行前向传播时，会进入 LeNet 的forward()函数，首先调用第一个Sequetial容器：self.features，由于self.features也是一个 module，因此会调用__call__()函数，里面调用result = self.forward(*input, **kwargs)，进入nn.Seuqetial的forward()函数，在这里依次调用所有的 module。

1
2
3
4
def forward(self, input):
        for module in self:
            input = module(input)
        return input

在上面可以看到在nn.Sequetial中，里面的每个子网络层 module 是使用序号来索引的，即使用数字来作为 key。一旦网络层增多，难以查找特定的网络层，这种情况可以使用 OrderDict (有序字典)。代码中使用

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
class LeNetSequentialOrderDict(nn.Module):
    def __init__(self, classes):
        super(LeNetSequentialOrderDict, self).__init__()

        self.features = nn.Sequential(OrderedDict({
            'conv1': nn.Conv2d(3, 6, 5),
            'relu1': nn.ReLU(inplace=True),
            'pool1': nn.MaxPool2d(kernel_size=2, stride=2),

            'conv2': nn.Conv2d(6, 16, 5),
            'relu2': nn.ReLU(inplace=True),
            'pool2': nn.MaxPool2d(kernel_size=2, stride=2),
        }))

        self.classifier = nn.Sequential(OrderedDict({
            'fc1': nn.Linear(16*5*5, 120),
            'relu3': nn.ReLU(),

            'fc2': nn.Linear(120, 84),
            'relu4': nn.ReLU(inplace=True),

            'fc3': nn.Linear(84, classes),
        }))
        ...
        ...
        ...

小结

nn.Sequetial是nn.Module的容器，用于按顺序包装一组网络层，有以下两个特性。

  • 顺序性：各网络层之间严格按照顺序构建，我们在构建网络时，一定要注意前后网络层之间输入和输出数据之间的形状是否匹配
  • 自带forward()函数：在nn.Sequetial的forward()函数里通过 for 循环依次读取每个网络层，执行前向传播运算。这使得我们我们构建的模型更加简洁

nn.ModuleList

nn.ModuleList是nn.Module的容器，用于包装一组网络层，以迭代的方式调用网络层。ModuleList 接收一个子模块（或层，需属于nn.Module类）的列表作为输入，然后也可以类似List那样进行append和extend操作。同时，子模块或层的权重也会自动添加到网络中来。主要有以下 3 个方法：

  • append()：在 ModuleList 后面添加网络层
  • extend()：拼接两个 ModuleList
  • insert()：在 ModuleList 的指定位置中插入网络层
1
2
3
4
net = nn.ModuleList([nn.Linear(784, 256), nn.ReLU()])
net.append(nn.Linear(256, 10)) # # 类似List的append操作
print(net[-1])  # 类似List的索引访问
print(net)
1
2
3
4
5
6
Linear(in_features=256, out_features=10, bias=True)
ModuleList(
  (0): Linear(in_features=784, out_features=256, bias=True)
  (1): ReLU()
  (2): Linear(in_features=256, out_features=10, bias=True)
)

下面的代码通过列表生成式来循环迭代创建 20 个全连接层，非常方便，只是在 forward()函数中需要手动调用每个网络层。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class ModuleList(nn.Module):
    def __init__(self):
        super(ModuleList, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(20)])

    def forward(self, x):
        for i, linear in enumerate(self.linears):
            x = linear(x)
        return x

net = ModuleList()

print(net)

fake_data = torch.ones((10, 10))

output = net(fake_data)

print(output)
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
ModuleList(
  (linears): ModuleList(
    (0-19): 20 x Linear(in_features=10, out_features=10, bias=True)
  )
)
tensor([[-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515],
        [-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515],
        [-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515],
        [-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515],
        [-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515],
        [-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515],
        [-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515],
        [-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515],
        [-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515],
        [-0.1253, -0.1615, -0.3969, -0.2678,  0.2628, -0.0124, -0.0215, -0.0008,
         -0.1870,  0.1515]], grad_fn=<AddmmBackward0>)

要特别注意的是，nn.ModuleList 并没有定义一个网络，它只是将不同的模块储存在一起。ModuleList中元素的先后顺序并不代表其在网络中的真实位置顺序，需要经过forward函数指定各个层的先后顺序后才算完成了模型的定义。具体实现时用for循环即可完成：

1
2
3
4
5
6
7
8
9
10
class model(nn.Module):
  def __init__(self, ...):
    super().__init__()
    self.modulelist = ...
    ...
    
  def forward(self, x):
    for layer in self.modulelist:
      x = layer(x)
    return x

nn.ModuleDict

nn.ModuleDict是nn.Module的容器，用于包装一组网络层，以索引的方式调用网络层，ModuleDict能够更方便地为神经网络的层添加名称。主要有以下 5 个方法：

  • clear()：清空 ModuleDict
  • items()：返回可迭代的键值对 (key, value)
  • keys()：返回字典的所有 key
  • values()：返回字典的所有 value
  • pop()：返回一对键值，并从字典中删除

下面的模型创建了两个ModuleDict：self.choices和self.activations，在前向传播时通过传入对应的 key 来执行对应的网络层。

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
class ModuleDict(nn.Module):
    def __init__(self):
        super(ModuleDict, self).__init__()
        self.choices = nn.ModuleDict({
            'conv': nn.Conv2d(10, 10, 3),
            'pool': nn.MaxPool2d(3)
        })

        self.activations = nn.ModuleDict({
            'relu': nn.ReLU(),
            'prelu': nn.PReLU()
        })

    def forward(self, x, choice, act):
        x = self.choices[choice](x)
        x = self.activations[act](x)
        return x

net = ModuleDict()

fake_img = torch.randn((4, 10, 32, 32))

output = net(fake_img, 'conv', 'relu')
# output = net(fake_img, 'conv', 'prelu')
print(output)
1
2
3
4
5
6
7
8
tensor([[[[0.0000, 0.0000, 0.0000,  ..., 0.4056, 0.1909, 0.0000],
          [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.9579, 0.4126],
          [0.0000, 0.6611, 0.0000,  ..., 0.0000, 0.3082, 0.3409],
          ...,
          [0.0000, 0.0000, 0.0000,  ..., 0.5356, 0.5810, 0.4790],
          [0.0000, 0.4047, 0.0221,  ..., 0.0910, 0.0000, 0.0000],
          [0.0000, 0.2876, 0.1203,  ..., 0.0000, 0.5455, 0.0000]]]],
       grad_fn=<ReluBackward0>)

容器总结

nn.Sequetial：顺序性，各网络层之间严格按照顺序执行，常用于 block 构建，在前向传播时的代码调用变得简洁

nn.ModuleList：迭代行，常用于大量重复网络构建，通过 for 循环实现重复构建

nn.ModuleDict：索引性，常用于可选择的网络层