PyTorch-网络的建立,预训练模型的加载

本文是PyTorch使用过程当中的的一些总结,有如下内容:python

  • 构建网络模型的方法
  • 网络层的遍历
  • 各层参数的遍历
  • 模型的保存与加载
  • 从预训练模型为网络参数赋值

主要涉及到如下函数的使用网络

  • add_module,ModulesList,Sequential 模型建立
  • modules(),named_modules(),children(),named_children() 访问模型的各个子模块
  • parameters(),named_parameters() 网络参数的遍历
  • save(),load()state_dict() 模型的保存与加载

构建网络

torch.nn.Module是全部网络的基类,在Pytorch实现的Model都要继承该类。并且,Module是能够包含其余的Module的,以树形的结构来表示一个网络结构。ide

简单的定义一个网络Model函数

class Model(nn.Module):
    def __init__(self):
        super(Model,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

Model中两个属性conv1conv2是两个卷积层,在正向传播的过程当中,再依次调用这两个卷积层。优化

除了使用Model的属性来为网络添加层外,还可使用add_module将网络层添加到网络中。google

class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.conv1 = nn.Conv2d(3,64,3)
        self.conv2 = nn.Conv2d(64,64,3)

        self.add_module("maxpool1",nn.MaxPool2d(2,2))
        self.add_module("covn3",nn.Conv2d(64,128,3))
        self.add_module("conv4",nn.Conv2d(128,128,3))

    def forward(self,x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.maxpool1(x)
        x = self.conv3(x)
        x = self.conv4(x)
        return x

add_module(name,layer)在正向传播的过程当中可使用添加时的name来访问改layer。url

这样一个个的添加layer,在简单的网络中还行,可是对于负责的网络层不少的网络来讲就须要敲不少重复的代码了。 这就须要使用到torch.nn.ModuleListtorch.nn.Sequentialcode

使用ModuleListSequential能够方便添加子网络到网络中,可是这二者仍是有所不一样的。orm

ModuleList

ModuleList是以list的形式保存sub-modules或者网络层,这样就能够先将网络须要的layer构建好保存到一个list,而后经过ModuleList方法添加到网络中。对象

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule,self).__init__()

        # 构建layer的list
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])

    def forward(self,x):

        # 正向传播,使用遍历每一个Layer
        for i, l in enumerate(self.linears):
            x = self.linears[i // 2](x) + l(x)

        return x

使用[nn.Linear(10, 10) for i in range(10)]构建要给Layer的list,而后使用ModuleList添加到网络中,在正向传播的过程当中,遍历该list

更为方便的是,能够提早配置后,所须要的各个Layer的属性,而后读取配置建立list,而后使用ModuleList将配置好的网络层添加到网络中。 以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加入到网络中。这样就能够很灵活的建立不一样的网络。

这里须要注意的是,ModuleList是将Module加入网络中,须要本身手动的遍历进行每个Moduleforward

Sequential

一个时序容器。Modules 会以他们传入的顺序被添加到容器中。固然,也能够传入一个OrderedDict一个时序容器。Modules 会以他们传入的顺序被添加到容器中。固然,也能够传入一个OrderedDict
Sequential也是一次加入多个Module到网络中中,和ModuleList不一样的是,它接受多个Module依次加入到网络中,还能够接受字典做为参数,例如:

# Example of using Sequential
        model = nn.Sequential(
                  nn.Conv2d(1,20,5),
                  nn.ReLU(),
                  nn.Conv2d(20,64,5),
                  nn.ReLU()
                )

# Example of using Sequential with OrderedDict
model = nn.Sequential(OrderedDict([
    ('conv1', nn.Conv2d(1,20,5)),
    ('relu1', nn.ReLU()),
    ('conv2', nn.Conv2d(20,64,5)),
    ('relu2', nn.ReLU())
    ]))

另外一个是,Sequential中实现了添加Module的forward,不须要手动的循环调用了。这点相比ModuleList较为方便。

总结

常见的有三种方法来添加子Module到网络中

  • 单独添加一个Module,可使用属性或者add_module方法。
  • ModuleList能够将一个Module的List加入到网络中,自由度较高,可是须要手动的遍历ModuleList进行forward
  • Sequential按照顺序将将Module加入到网络中,也能够处理字典。 相比于ModuleList不须要本身实现forward

遍历网络结构

可使用如下2对4个方法来访问网络层全部的Modules

  • modules()named_modules()
  • children()named_children()

modules方法

简单的定义一个以下网络:

class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3,out_channels=64,kernel_size=3)
        self.conv2 = nn.Conv2d(64,64,3)
        self.maxpool1 = nn.MaxPool2d(2,2)

        self.features = nn.Sequential(OrderedDict([
            ('conv3', nn.Conv2d(64,128,3)),
            ('conv4', nn.Conv2d(128,128,3)),
            ('relu1', nn.ReLU())
        ]))

    def forward(self,x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.maxpool1(x)
        x = self.features(x)

        return x

modules()方法,返回一个包含当前模型全部模块的迭代器,这个是递归的返回网络中的全部Module。使用以下语句

m = Model()
    for idx,m in enumerate(m.modules()):
        print(idx,"-",m)

其结果为:

0 - Model(
  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
  (maxpool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (features): Sequential(
    (conv3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))
    (conv4): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
    (relu1): ReLU()
  )
)
1 - Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
2 - Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
3 - MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
4 - Sequential(
  (conv3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))
  (conv4): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
  (relu1): ReLU()
)
5 - Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))
6 - Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
7 - ReLU()

输出结果解析:

  • 0-Model 整个网络模块
  • 1-2-3-4 为网络的4个子模块,注意4 - Sequential仍然包含有子模块
  • 5-6-7为模块4 - Sequential的子模块

能够看出modules()是递归的返回网络的各个module,从最顶层直到最后的叶子module。

named_modules()的功能和modules()的功能相似,不一样的是它返回内容有两部分:module的名称以及module。

children()方法

modules()不一样,children()只返回当前模块的子模块,不会递归子模块。

for idx,m in enumerate(m.children()):
        print(idx,"-",m)

其输出为:

0 - Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
1 - Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
2 - MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
3 - Sequential(
  (conv3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))
  (conv4): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
  (relu1): ReLU()
)

子模块3-Sequential仍然有子模块,children()没有递归的返回。
named_children()children()的功能相似,不一样的是其返回两部份内容:模块的名称以及模块自己。

网络的参数

方法parameters()返回一个包含模型全部参数的迭代器。通常用来看成optimizer的参数。

for p in m.parameters():
        print(type(p.data),p.size())

其输出为:

<class 'torch.Tensor'> torch.Size([128, 64, 3, 3])
<class 'torch.Tensor'> torch.Size([128])
<class 'torch.Tensor'> torch.Size([128, 128, 3, 3])
<class 'torch.Tensor'> torch.Size([128])

包含网络中的全部的权值矩阵参数以及偏置参数。 对网络进行训练时须要将parameters()做为优化器optimizer的参数。

optimizer = torch.optim.SGD(m1.parameters(),lr = args.lr,momentum=args.momentum,weight_decay=args.weight_decay)

parameters()返回网络的全部参数,主要是提供给optimizer用的。而要取得网络某一层的参数或者参数进行一些特殊的处理(如fine-tuning),则使用named_parameters()更为方便些。

named_parameters()返回参数的名称及参数自己,能够按照参数名对一些参数进行处理。

以上面的vgg网络为例:

for k,v in m1.named_parameters():
    print(k,v.size())

named_parameters返回的是键值对,k为参数的名称 ,v为参数自己。输出结果为:

vgg.0.weight torch.Size([64, 3, 3, 3])
vgg.0.bias torch.Size([64])
vgg.2.weight torch.Size([64, 64, 3, 3])
vgg.2.bias torch.Size([64])
vgg.5.weight torch.Size([128, 64, 3, 3])
vgg.5.bias torch.Size([128])
vgg.7.weight torch.Size([128, 128, 3, 3])
vgg.7.bias torch.Size([128])
vgg.10.weight torch.Size([256, 128, 3, 3])
vgg.10.bias torch.Size([256])
vgg.12.weight torch.Size([256, 256, 3, 3])
vgg.12.bias torch.Size([256])
vgg.14.weight torch.Size([256, 256, 3, 3])
vgg.14.bias torch.Size([256])
vgg.17.weight torch.Size([512, 256, 3, 3])
vgg.17.bias torch.Size([512])
vgg.19.weight torch.Size([512, 512, 3, 3])
vgg.19.bias torch.Size([512])
vgg.21.weight torch.Size([512, 512, 3, 3])
vgg.21.bias torch.Size([512])
vgg.24.weight torch.Size([512, 512, 3, 3])
vgg.24.bias torch.Size([512])
vgg.26.weight torch.Size([512, 512, 3, 3])
vgg.26.bias torch.Size([512])
vgg.28.weight torch.Size([512, 512, 3, 3])
vgg.28.bias torch.Size([512])

参数名的命名规则属性名称.参数属于的层的编号.weight/bias。 这在fine-tuning的时候,给一些特定的层的参数赋值是很是方便的,这点在后面在加载预训练模型时会看到。

模型的保存与加载

PyTorch使用torch.savetorch.load方法来保存和加载网络,并且网络结构和参数能够分开的保存和加载。

  • 保存网络结构及其参数
torch.save(model,'model.pth') # 保存
model = torch.load("model.pth") # 加载
  • 只加载模型参数,网络结构从代码中建立
torch.save(model.state_dict(),"model.pth") # 保存参数
model = model() # 代码中建立网络结构
params = torch.load("model.pth") # 加载参数
model.load_state_dict(params) # 应用到网络结构中

加载预训练模型

PyTorch中的torchvision里有不少经常使用网络的预训练模型,例如:vgg,resnet,googlenet等,能够方便的使用这些预训练模型进行微调。

# PyTorch中的torchvision里有不少经常使用的模型,能够直接调用:
import torchvision.models as models
 
resnet101 = models.resnet18(pretrained=True)
alexnet = models.alexnet()
squeezenet = models.squeezenet1_0()

有时候只须要加载预训练模型的部分参数,可使用参数名做为过滤条件,以下

resnet152 = models.resnet152(pretrained=True)
pretrained_dict = resnet152.state_dict()
"""加载torchvision中的预训练模型和参数后经过state_dict()方法提取参数
   也能够直接从官方model_zoo下载:
   pretrained_dict = model_zoo.load_url(model_urls['resnet152'])"""
model_dict = model.state_dict()
# 将pretrained_dict里不属于model_dict的键剔除掉
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 更新现有的model_dict
model_dict.update(pretrained_dict)
# 加载咱们真正须要的state_dict
model.load_state_dict(model_dict)

model.state_dict()返回一个python的字典对象,将每一层与它的对应参数创建映射关系(如model的每一层的weights及偏置等等)。注意,只有有参数训练的层才会被保存。

上述的加载方式,是按照参数名类匹配过滤的,可是对于一些参数名称没法彻底匹配,或者在预训练模型的基础上新添加的一些层,这些层没法从预训练模型中获取参数,须要初始化。

仍然以上述的vgg为例,在标准的vgg16的特征提取后面,新添加两个卷积层,这两个卷积层的参数须要进行初始化。

vgg = torch.load("vgg.pth") # 加载预训练模型

for k,v in m1.vgg.named_parameters():
    k = "features.{}".format(k) # 参数名称
    if k in vgg.keys():
        v.data = vgg[k].data # 直接加载预训练参数
    else:
        if k.find("weight") >= 0:
            nn.init.xavier_normal_(v.data) # 没有预训练,则使用xavier初始化
        else:
            nn.init.constant_(v.data,0) # bias 初始化为0
相关文章
相关标签/搜索