本篇博客讲述了如何使用 nn.ModuleList() 和 nn.Sequential() 简化模型的创建方式。并分别使用传统方法,nn.ModuleList() 以及 nn.Sequential() 创建一个 拟合 sin 函数的模型,训练模型保存结果,体会两个函数的方便之处。
一、使用传统方法
创建模型
import numpy as np
import torch
import torch.nn as nn#准备数据
data=np.linspace(-2*np.pi,2*np.pi,400) # -2PI 到 2PI 区间分成400份的所有点
x=torch.tensor(data.reshape(400,-1),dtype=torch.float)
y=torch.tensor(np.sin(data.reshape(400,-1)),dtype=torch.float)# 创建模型
class model1(nn.Module):def __init__(self):super(model1,self).__init__()self.linear1=nn.Linear(1,10)self.activation1=nn.ReLU()self.linear2=nn.Linear(10,100)self.activation2=nn.ReLU()self.linear3=nn.Linear(100,10)self.activation3=nn.ReLU()self.linear4=nn.Linear(10,1)def forward(self,x):out=self.linear1(x)out=self.activation1(out)out=self.linear2(out)out=self.activation2(out)out=self.linear3(out)out=self.activation3(out)out=self.linear4(out)return out# 创建模型实例
model=model1()
## 设置优化器和损失函数
optimizer=torch.optim.Adam(model1.parameters(model),lr=0.005)
Loss=nn.MSELoss()# 开始训练
for i in range(3000):y_predict=model(x)loss=Loss(y_predict,y)optimizer.zero_grad()loss.backward()optimizer.step()if (i+1)%100==0:print("[training] step: {0} , loss: {1}".format(i+1,loss))## 展示和保存图片
import matplotlib.pyplot as plt
predict=model(x)
plt.plot(data,np.sin(data),color="red")
plt.plot(data,predict.detach().numpy(),color="blue")
plt.savefig(figsize=(10,10),fname="result.png")
plt.show()
输出结果:
红色为训练的数据集(真实的数据),蓝色为预测的结果。
可以看到训练了 3000次之后,拟合的还不错
二、使用 nn.ModuleList() 函数
传统的创建模型方法代码如下所示:
class model1(nn.Module):def __init__(self):super(model1,self).__init__()self.linear1=nn.Linear(1,10)self.activation1=nn.ReLU()self.linear2=nn.Linear(10,100)self.activation2=nn.ReLU()self.linear3=nn.Linear(100,10)self.activation3=nn.ReLU()self.linear4=nn.Linear(10,1)def forward(self,x):out=self.linear1(x)out=self.activation1(out)out=self.linear2(out)out=self.activation2(out)out=self.linear3(out)out=self.activation3(out)out=self.linear4(out)return out
使用 ModuleList 可以简化写法。
这里需要讲的是,ModuleList 可以存储多个 model,传统的方法,一个model 就要写一个 forward ,但是如果将它们存到一个 ModuleList 的话,就可以使用一个 forward。
ModuleList是Module的子类,当在Module中使用它的时候,就能自动识别为子module。
当添加 nn.ModuleList 作为 nn.Module 对象的一个成员时(即当我们添加模块到我们的网络时),所有 nn.ModuleList 内部的 nn.Module 的 parameter 也被添加作为 我们的网络的 parameter。
使用 ModuleList 也可以使得网络的结构具有灵活性,比如我需要将网络的层数设置为变量,传统的方法要借助 list 实现,并且不方便,而使用 ModuleList就可以简化这个操作。
使用ModuleList替换原来的方法,修改之后的模型如下:(输入数据,优化器,损失,训练过程等不便,这里省略)
class model2(nn.Module):def __init__(self):super(model2, self).__init__()self.layers=nn.ModuleList([nn.Linear(1,10), nn.ReLU(),nn.Linear(10,100),nn.ReLU(),nn.Linear(100,10),nn.ReLU(),nn.Linear(10,1)])def forward(self,x):out=xfor i,layer in enumerate(self.layers):out=layer(out)return out
其它用法
ModuleList 具有和List 相似的用法,实际上可以把它视作是 Module 和 list 的结合。
除了在创建 ModuleList 的时候传入一个 module 的 列表,还可以使用extend 函数和 append 函数来添加模型
- extend 方法
和 list 相似,参数为一个元素为 Module的列表,该方法的效果是将列表中的所有 Module 添加到 ModuleList中。self.linears.extend([nn.Linear(size1, size2) for i in range(1, num_layers)])
- append 方法
和list 的append 方法一样,将 一个 Module 添加到ModuleList。self.linears.append(nn.Linear(size1, size2)
三、使用 nn.Sequential()
修改之后的模型如下
class model3(nn.Module):def __init__(self):super(model3, self).__init__()self.network=nn.Sequential(nn.Linear(1,10),nn.ReLU(),nn.Linear(10,100),nn.ReLU(),nn.Linear(100,10),nn.ReLU(),nn.Linear(10,1))def forward(self, x):return self.network(x)
可以使用 self.network[0] 获取第一个 Linear子模型,由于每一个子模型没有设置唯一的名称,所以只能使用数字索引来获取。
添加子模型的方法
与ModuleList 一样,Sequential 方法除了在创建的时候直接加入子模型,也可以在创建之后加入新的模型,而且可以给子模型定义唯一的名称索引,方便获取,示例如下:
self.network.add_module("linear1",nn.Linear(100,100))
如果想获得这个子模型,只要如下操作即可:
linear=self.network.linear1
另外一种写法
from collections import OrderedDict
self.network=nn.Sequential(OrderedDict(("linear1",nn.Linear(1,10)),("activation1",nn.ReLU()),("linear2",nn.Linear(10,100))))
同样,获取子模型可以使用名称获取,与第二种写法一致。
完整参考代码
完整参考代码如下,修改标注的位置可以切换模型的不同搭建方法
import numpy as np
import torch
import torch.nn as nn#准备数据
data=np.linspace(-2*np.pi,2*np.pi,400) # -2PI 到 2PI 区间分成400份的所有点
x=torch.tensor(data.reshape(400,-1),dtype=torch.float)
y=torch.tensor(np.sin(data.reshape(400,-1)),dtype=torch.float)# 创建模型
class model1(nn.Module):"""第一种搭建方式"""def __init__(self):super(model1,self).__init__()self.linear1=nn.Linear(1,10)self.activation1=nn.ReLU()self.linear2=nn.Linear(10,100)self.activation2=nn.ReLU()self.linear3=nn.Linear(100,10)self.activation3=nn.ReLU()self.linear4=nn.Linear(10,1)def forward(self,x):out=self.linear1(x)out=self.activation1(out)out=self.linear2(out)out=self.activation2(out)out=self.linear3(out)out=self.activation3(out)out=self.linear4(out)return out
class model2(nn.Module):"""第二种搭建方式"""def __init__(self):super(model2, self).__init__()self.layers=nn.ModuleList([nn.Linear(1,10),nn.ReLU(),nn.Linear(10,100),nn.ReLU(),nn.Linear(100,10),nn.ReLU(),nn.Linear(10,1)])def forward(self,x):out=xfor i,layer in enumerate(self.layers):out=layer(out)return outdef printf(self):for i,layer in enumerate(self.layers):print("第{0}个层:{1}".format(i,layer))class model3(nn.Module):"""第三种搭建方式"""def __init__(self):super(model3, self).__init__()self.network=nn.Sequential(nn.Linear(1,10),nn.ReLU(),nn.Linear(10,100),nn.ReLU(),nn.Linear(100,10),nn.ReLU(),nn.Linear(10,1))def forward(self, x):return self.network(x)
# 创建模型实例
#####################################################################################################
################################ 这里切换模型的搭建方式 ############################################
#####################################################################################################
model=model3()## 设置优化器和损失函数
optimizer=torch.optim.Adam(model1.parameters(model),lr=0.005)
Loss=nn.MSELoss()# 开始训练
for i in range(3000):y_predict=model(x)loss=Loss(y_predict,y)optimizer.zero_grad()loss.backward()optimizer.step()if (i+1)%100==0:print("[training] step: {0} , loss: {1}".format(i+1,loss))## 展示和保存图片
import matplotlib.pyplot as plt
predict=model(x)
plt.plot(data,np.sin(data),color="red")
plt.plot(data,predict.detach().numpy(),color="blue")
plt.savefig(figsize=(10,10),fname="result.png")
plt.show()