本文最后更新于:2024年5月7日 下午

这里我们来更为系统地学习PyTorch中模型定义的方式,本节的学习将为后续灵活构建自己的模型打下坚实的基础。

参考 深入浅出PyTorch ,系统补齐基础知识。

本节目录

  • 熟悉PyTorch中模型定义的三种方式
  • 读懂GitHub上千奇百怪的写法
  • 自己根据需要灵活选取模型定义方式

简介

Module 类是 torch.nn 模块里提供的一个模型构造类 (nn.Module),是所有神经⽹网络模块的基类,我们可以继承它来定义我们想要的模型;

PyTorch模型定义应包括两个主要部分:各个部分的初始化(__init__);数据流向定义(forward

基于nn.Module,我们可以通过SequentialModuleListModuleDict三种方式定义PyTorch模型。

Sequential

对应模块为nn.Sequential()

当模型的前向计算为简单串联各个层的计算时, Sequential 类可以通过更加简单的方式定义模型。它可以接收一个子模块的有序字典(OrderedDict) 或者一系列子模块作为参数来逐一添加 Module 的实例,⽽模型的前向计算就是将这些实例按添加的顺序逐⼀计算。我们结合Sequential和定义方式加以理解:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
from collections import OrderedDict
class MySequential(nn.Module):
def __init__(self, *args):
super(MySequential, self).__init__()
if len(args) == 1 and isinstance(args[0], OrderedDict): # 如果传入的是一个OrderedDict
for key, module in args[0].items():
self.add_module(key, module)
# add_module方法会将module添加进self._modules(一个OrderedDict)
else: # 传入的是一些Module
for idx, module in enumerate(args):
self.add_module(str(idx), module)
def forward(self, input):
# self._modules返回一个 OrderedDict,保证会按照成员添加时的顺序遍历成
for module in self._modules.values():
input = module(input)
return input

下面来看下如何使用Sequential来定义模型。只需要将模型的层按序排列起来即可,根据层名的不同,排列的时候有两种方式:

  • 直接排列
1
2
3
4
5
6
7
8
9
10
11
12
import torch.nn as nn
net = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 10),
)
print(net)
Sequential(
(0): Linear(in_features=784, out_features=256, bias=True)
(1): ReLU()
(2): Linear(in_features=256, out_features=10, bias=True)
)
  • 使用OrderedDict:
1
2
3
4
5
6
7
8
9
10
11
12
13
import collections
import torch.nn as nn
net2 = nn.Sequential(collections.OrderedDict([
('fc1', nn.Linear(784, 256)),
('relu1', nn.ReLU()),
('fc2', nn.Linear(256, 10))
]))
print(net2)
Sequential(
(fc1): Linear(in_features=784, out_features=256, bias=True)
(relu1): ReLU()
(fc2): Linear(in_features=256, out_features=10, bias=True)
)

可以看到,使用Sequential定义模型的好处在于简单、易读,同时使用Sequential定义的模型不需要再写forward,因为顺序已经定义好了。但使用Sequential也会使得模型定义丧失灵活性,比如需要在模型中间加入一个外部输入时就不适合用Sequential的方式实现。使用时需根据实际需求加以选择。

ModuleList

对应模块为nn.ModuleList()

ModuleList 接收一个子模块(或层,需属于 nn.Module 类)的列表作为输入,然后也可以类似List那样进行 append 和 extend 操作。同时,子模块或层的权重也会自动添加到网络中来。

1
2
3
4
5
6
7
8
9
10
net = nn.ModuleList([nn.Linear(784, 256), nn.ReLU()])
net.append(nn.Linear(256, 10)) # # 类似List的append操作
print(net[-1]) # 类似List的索引访问
print(net)
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)
)

要特别注意的是,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

ModuleDict

对应模块为nn.ModuleDict()

ModuleDictModuleList的作用类似,只是ModuleDict能够更方便地为神经网络的层添加名称。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
net = nn.ModuleDict({
'linear': nn.Linear(784, 256),
'act': nn.ReLU(),
})
net['output'] = nn.Linear(256, 10) # 添加
print(net['linear']) # 访问
print(net.output)
print(net)
Linear(in_features=784, out_features=256, bias=True)
Linear(in_features=256, out_features=10, bias=True)
ModuleDict(
(act): ReLU()
(linear): Linear(in_features=784, out_features=256, bias=True)
(output): Linear(in_features=256, out_features=10, bias=True)
)

三种方法的比较与适用场景

Sequential适用于快速验证结果,因为已经明确了要用哪些层,直接写一下就好了,不需要同时写__init__forward

ModuleListModuleDict 在某个完全相同的层需要重复出现多次时,非常方便实现,可以”一行顶多行“;

当我们需要之前层的信息的时候,比如 ResNets 中的残差计算,当前层的结果需要和之前层中的结果进行融合,一般使用 ModuleList/ModuleDict 比较方便。

参考资料



文章链接:
https://www.zywvvd.com/notes/study/deep-learning/pytorch/torch-learning/torch-learning-9/


“觉得不错的话,给点打赏吧 ୧(๑•̀⌄•́๑)૭”

微信二维码

微信支付

支付宝二维码

支付宝支付

PyTorch 学习 -9- 模型定义
https://www.zywvvd.com/notes/study/deep-learning/pytorch/torch-learning/torch-learning-9/
作者
Yiwei Zhang
发布于
2023年7月23日
许可协议