0%

[DL] PyTorch 折桂 6：torch.nn.Module

发表于 更新于

本文中，我们看一看如何构建模型。
创造一个模型分两步：构建模型和权值初始化。而构建模型又有“定义单独的网络层”和“把它们拼在一起”两步。

1. torch.nn.Module

torch.nn.Module 是所有 torch.nn 中的类的父类。我们来看一个非常简单的神经网络：

1
2
3
4
5
6
7
8
class SimpleNet(nn.Module):
    def __init__(self, x):
        super(SimpleNet,self).__init__()
        self.fc = nn.Linear(x.shape[0], 1)
        
    def forward(self, x):
        x = self.fc(x)
        return x

我们随便喂给它一个张量，打印它的网络：

1
2
3
4
5
>>> simpleNet = SimpleNet(torch.tensor((10, 2)))
>>> print(simpleNet)
SimpleNet(
  (fc): Linear(in_features=2, out_features=1, bias=True)
)

所有自定义的神经网络都要继承 torch.nn.Module。定义单独的网络层在 __init__ 函数中实现，把定义好的网络层拼接在一起在 forward 函数中实现。网络类有两个重要的函数：parameters 存储了模型的权重；modules 存储了模型的结构。

1
2
3
4
5
6
7
8
9
10
11
>>> list(simpleNet.modules())
[SimpleNet(
   (fc): Linear(in_features=2, out_features=1, bias=True)
 ),
 Linear(in_features=2, out_features=1, bias=True)]
 
 >>> list(simpleNet.parameters())
[Parameter containing:
 tensor([[ 0.1533, -0.2574]], requires_grad=True),
 Parameter containing:
 tensor([-0.1589], requires_grad=True)]

2. torch.nn.Sequential

这是一个序列容器，既可以放在模型外面单独构建一个模型，也可以放在模型里面成为模型的一部分。

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
27
28
29
30
31
# 单独成为一个模型
model1 = nn.Sequential(
          nn.Conv2d(1,20,5),
          nn.ReLU(),
          nn.Conv2d(20,64,5),
          nn.ReLU()
        )
# 成为模型的一部分
class LeNetSequential(nn.Module):
    def __init__(self, classes):
        super(LeNetSequential, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=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

放在模型里面的话，模型还是需要 __init__ 和 forward 函数。

这样构建出来的模型的层没有名字：

1
2
3
4
5
6
7
8
9
10
11
12
13
>>> model2 = nn.Sequential(
...           nn.Conv2d(1,20,5),
...           nn.ReLU(),
...           nn.Conv2d(20,64,5),
...           nn.ReLU()
...         )
>>> model2
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()
)

为了方便区分不同的层，我们可以使用 collections 里的 OrderedDict 函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
>>> from collections import OrderedDict
>>> model3 = nn.Sequential(OrderedDict([
...           ('conv1', nn.Conv2d(1,20,5)),
...           ('relu1', nn.ReLU()),
...           ('conv2', nn.Conv2d(20,64,5)),
...           ('relu2', nn.ReLU())
...         ]))
>>> model3
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()
)

3. torch.nn.ModuleList

将网络层存储进一个列表，可以使用列表生成式快速生成网络，生成的网络层可以被索引，也拥有列表的方法 append，extend 或 insert。

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
27
>>> class MyModule(nn.Module):
...     def __init__(self):
...         super(MyModule, self).__init__()
...         self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])
...         self.linears.append(nn.Linear(10, 1)) # append
...     def forward(self, x):
...         for i, l in enumerate(self.linears):
...             x = self.linears[i // 2](x) + l(x)
...         return x
    
>>> myModeul = MyModule()
>>> myModeul
MyModule(
  (linears): ModuleList(
    (0): Linear(in_features=10, out_features=10, bias=True)
    (1): Linear(in_features=10, out_features=10, bias=True)
    (2): Linear(in_features=10, out_features=10, bias=True)
    (3): Linear(in_features=10, out_features=10, bias=True)
    (4): Linear(in_features=10, out_features=10, bias=True)
    (5): Linear(in_features=10, out_features=10, bias=True)
    (6): Linear(in_features=10, out_features=10, bias=True)
    (7): Linear(in_features=10, out_features=10, bias=True)
    (8): Linear(in_features=10, out_features=10, bias=True)
    (9): Linear(in_features=10, out_features=10, bias=True)
    (10): Linear(in_features=10, out_features=1, bias=True) # append 进的层
  )
)

4. torch.nn.ModuleDict

这个函数与上面的 torch.nn.Sequential(OrderedDict(...)) 的行为非常类似，并且拥有 keys，values，items，pop，update 等词典的方法：

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
27
>>> class MyDictDense(nn.Module):
...     def __init__(self):
...         super(MyDictDense, self).__init__()
...         self.params = nn.ModuleDict({
...                 'linear1': nn.Linear(512, 128),
...                 'linear2': nn.Linear(128, 32)
...         })
...         self.params.update({'linear3': nn.Linear(32, 10)}) # 添加层

...     def forward(self, x, choice='linear1'):
...         return torch.mm(x, self.params[choice])

>>> net = MyDictDense()
>>> print(net)
MyDictDense(
  (params): ModuleDict(
    (linear1): Linear(in_features=512, out_features=128, bias=True)
    (linear2): Linear(in_features=128, out_features=32, bias=True)
    (linear3): Linear(in_features=32, out_features=10, bias=True)
  )
)

>>> print(net.params.keys())
odict_keys(['linear1', 'linear2', 'linear3'])

>>> print(net.params.items())
odict_items([('linear1', Linear(in_features=512, out_features=128, bias=True)), ('linear2', Linear(in_features=128, out_features=32, bias=True)), ('linear3', Linear(in_features=32, out_features=10, bias=True))])

欢迎关注我的微信公众号“花解语 NLP”：
在这里插入图片描述

欢迎关注我的其它发布渠道