一、继承Module类来构造模型
Module类是nn模块里提供的一个模型构造类,是所有神经网络模块的基类,如下所示,继承这个类需要重载Module类中的__init__
函数和forward
函数,它们分别用于创建模型参数和定义前向计算(正向传播)。
继承时无需定义反向传播函数,因为系统将通过自动求梯度而自动生成反向传播所需的backward
函数。
如下的代码定义了一个MLP类:
import torch
from torch import nn
class MLP(nn.Module):
def __init__(self,**kwargs):
super(MLP,self).__init__(**kwargs)
self.hidden=nn.Linear(784,256)
self.act=nn.ReLU()
self.output=nn.Linear(256,10)
def forward(self,x):
a=self.act(self.hidden(x))
return self.output(a)
我们可以实例化MLP类得到模型变量net。下面的代码初始化net并传入输入数据x进行一次前向计算。其中,net(X)会调用MLP继承自Module类的__call__函数,这个函数将调用MLP类定义的forward函数来完成前向计算。
x=torch.rand(2,784)
net=MLP()
print(net)
net(X)
输出:
MLP(
(hidden): Linear(in_features=784, out_features=256, bias=True)
(act): ReLU()
(output): Linear(in_features=256, out_features=10, bias=True)
)
二、Module的子类
Module类是一个通用的部件,Pytorch还实现了继承自Module的可以方便构建模型的类,如Sequential、ModuleList和ModuleDict等等。
2.1 Sequential类
当模型的前向计算为简单串联各个层的计算时, Sequential 类可以通过更加简单的⽅式定义模型。这正是 Sequential 类的⽬的:它可以接收⼀个⼦模块的有序字典(OrderedDict)或者⼀系列⼦模块作为参数来逐⼀添加 Module 的实例,⽽模型的前向计算就是将这些实例按添加的顺序逐⼀计算。
2.2 ModuleList类
ModuleList接收一个子模块的列表作为输入,然后也可以类似List那样进行append和extend操作:
net=nn.ModuleList([nn.Linear(784,256),nn.ReLU()])
net.append(nn.Linear(256,10))
print(net[-1])
print(net)
输出:
ModuleList(
(0): Linear(in_features=784, out_features=256, bias=True)
(1): ReLU()
(2): Linear(in_features=256, out_features=10, bias=True)
)
2.3 ModuleDict类
ModuleDict接收一个子模块的字典作为输入,然后也可以类似字典那样进行添加访问操作:
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(
(linear): Linear(in_features=784, out_features=256, bias=True)
(act): ReLU()
(output): Linear(in_features=256, out_features=10, bias=True)
)
三、构造复杂的模型
虽然上面介绍的这些类可以使模型构造更加简单,且不需要定义forward函数,但直接继承Module类可以极大地拓展模型构造的灵活性。下面我们创建一个稍微复杂些的网络FancyMLP。在这个网络中,我们通过get_constant函数创建训练中不被迭代的参数,即常数参数。在前向计算中,除了使用创建的常数参数外,我们还使用Tensor的函数和Python的控制流,并多次调用相同的层。
class FancyMLP(nn.Module):
def __init__(self,**kwargs):
super(FancyMLP,self).__init__(**kwargs)
#这里定义了不可训练的参数(常数参数)
self.rand_weight=torch.rand((20,20),requires_grad=False)
self.linear=nn.Linear(20,20)
def forward(self,x):
x=self.linear(x)
#使用创建的常数参数,以及nn.functional中的relu函数以及mm函数
x=nn.functional.relu(torch.mm(x,self.rand_weight.data)+1)
#复用全连接层,等价于两个全连接层共享参数
x=self.linear(x)
#控制流,这里我们需要调用item函数来返回标量进行比较
while x.norm().item()>1:
x/=2
if x.norm().item()<0.8:
x*=10
return x.sum()
测试:
X=torch.rand(2,20)
net=FancyMLP()
print(net)
print(net(X))
输出:
FancyMLP(
(linear): Linear(in_features=20, out_features=20, bias=True)
)
tensor(1.5884, grad_fn=<SumBackward0>)
因为FancyMLP和Sequential类都是Module类的子类,所以我们可以嵌套调用它们。
class NestMLP(nn.Module):
def __init__(self,**kwargs):
super(NestMLP,self).__init__(**kwargs)
self.net=nn.Sequential(nn.Linear(40,30),nn.ReLU())
def forward(self,x):
return self.net(x)
net=nn.Sequential(NestMLP(),nn.Linear(30,20),FancyMLP())
X=torch.rand(2,40)
print(net)
print(net(X))
输出:
Sequential(
(0): NestMLP(
(net): Sequential(
(0): Linear(in_features=40, out_features=30, bias=True)
(1): ReLU()
)
)
(1): Linear(in_features=30, out_features=20, bias=True)
(2): FancyMLP(
(linear): Linear(in_features=20, out_features=20, bias=True)
)
)
tensor(4.6094, grad_fn=<SumBackward0>)