本文是对一些函数的学习。函数主要包括下面四个方便:
- 模型构建的函数:
add_module,add_module,add_module - 访问子模块:
add_module,add_module,add_module,add_module - 网络遍历:
add_module,add_module - 模型的保存与加载:
add_module,add_module,add_module
1 模型构建函数
torch.nn.Module是所有网络的基类,在PyTorch实现模型的类中都要继承这个类(这个在之前的课程中已经提到)。在构建Module中,Module是一个包含其他的Module的,类似于,你可以先定义一个小的网络模块,然后把这个小模块作为另外一个网络的组件。因此网络结构是呈现树状结构。
我们先简单定义一个网络:
import torch.nn as nn
import torch
class MyNet(nn.Module):
def __init__(self):
super(MyNet,self).__init__()
self.conv1 = nn.Conv2d(3,64,3)
self.conv2 = nn.Conv2d(64,64,3)
def forward(self,x):
x = self.conv1(x)
x = self.conv2(x)
return x
net = MyNet()
print(net)
输出结果:
MyNet中有两个属性conv1和conv2是两个卷积层,在正向传播forward的过程中,依次调用这两个卷积层实现网络的功能。
1.1 add_module
这种是最常见的定义网络的功能,在有些项目中,会看到这样的方法add_module。我们用这个方法来重写上面的网络:
class MyNet(nn.Module):
def __init__(self):
super(MyNet,self).__init__()
self.add_module('conv1',nn.Conv2d(3,64,3))
self.add_module('conv2',nn.Conv2d(64,64,3))
def forward(self,x):
x = self.conv1(x)
x = self.conv2(x)
return x
其实add_module(name,layer)和self.name=layer实现了相同的功能,个人感觉也许是因为add_module可以使用字符串来定义变量名字,所以可以放在循环中?反正这个先了解熟悉熟悉。
上面的两种方法都是一层一层的添加layer,如果网络复杂的话,那就需要写很多重复的代码了。因此接下来来讲解一下网络模块的构建,torch.nn.ModuleList和torch.nn.Sequential
1.2 ModuleList
ModuleList按照字面意思是用list的形式保存网络层的。这样就可以先将网络需要的layer构建好,保存到一个list,然后通过ModuleList方法添加到网络中.
class MyNet(nn.Module):
def __init__(self):
super(MyNet,self).__init__()
self.linears = nn.ModuleList(
[nn.Linear(10,10) for i in range(5)]
)
def forward(self,x):
for l in self.linears:
x = l(x)
return x
net = MyNet()
print(net)
输出结果是:
这个ModuleList主要是用在读取config文件来构建网络模型中的,下面用VGG模型的构建为例子:
vgg_cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
512, 512, 512, 'M']
def vgg(cfg, i, batch_norm=False):
layers = []
in_channels = i
for v in cfg:
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
elif v == 'C':
layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
if batch_norm:
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
else:
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
return layers
class Model1(nn.Module):
def __init__(self):
super(Model1,self).__init__()
self.vgg = nn.ModuleList(vgg(vgg_cfg,3))
def forward(self,x):
for l in self.vgg:
x = l(x)
m1 = Model1()
print(m1)
先读取网络结构的配置文件vgg_cfg然后根据这个文件创建对应的Layer list,然后使用ModuleList添加到网络中,这样可以快速创建不同的网络(用上面为例子的话,可以通过修改配置文件,然后快速修改网络结构 )
1.3 Sequential
在一些自己做的小项目中,Sequential其实用的更为频繁。
依然重写最初最简单的例子:
class MyNet(nn.Module):
def __init__(self):
super(MyNet,self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(3,64,3),
nn.Conv2d(64,64,3)
)
def forward(self,x):
x = self.conv(x)
return x
net = MyNet()
print(net)
运行结果:
观察细致的朋友可以发现这个问题,Seqential内的网络层是默认用数字进行标号的,而一开始我们使用self.conv1和self.conv2的时候,使用conv1和conv2作为标号的。
我们如何修改Sequential中网络层的名称呢?这里需要使用到collections.OrderedDict有序字典。Sequential是支持有序字典构建的。
from collections import OrderedDict
class MyNet(nn.Module):
def __init__(self):
super(MyNet,self).__init__()
self.conv = nn.Sequential(OrderedDict([
('conv1',nn.Conv2d(3,64,3)),
('conv2',nn.Conv2d(64,64,3))
]))
def forward(self,x):
x = self.conv(x)
return x
net = MyNet()
print(net)
输出结果:
1.4 小总结
- 单独增加一个网络层或者子模块,可以用
add_module或者直接赋予属性; ModuleList可以将一个Module的List增加到网络中,自由度较高。Sequential按照顺序产生一个Module模块。这里推荐习惯使用OrderedDict的方法进行构建。对网络层加上规范的名称,这样有助于后续查找与遍历