Pytorch模型构造方法

一、继承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>)