PyTorch-超全的网络创建，显示，预训练模型的加载

最新推荐文章于 2023-04-27 20:58:09 发布

ChaoFeiLi

最新推荐文章于 2023-04-27 20:58:09 发布

阅读量646

点赞数 3

分类专栏： Pytorch 文章标签： pytorch

原文链接：https://www.cnblogs.com/wangguchangqing/p/11058525.html

版权

Pytorch 专栏收录该内容

27 篇文章 11 订阅

订阅专栏

参考链接：https://www.cnblogs.com/wangguchangqing/p/11058525.html

本文是PyTorch使用过程中的的一些总结，有以下内容：

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

主要涉及到以下函数的使用：

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

1、构建网络

torch.nn.Module是所有网络的基类，在Pytorch实现的Model都要继承该类。而且，Module是可以包含其他的Module的，以树形的结构来表示一个网络结构。

这一个小段落将介绍三个创建网络的方法：add_module，ModulesList，Sequential

add_module：

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。

使用add_module定义一个类，然后构建模型

import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

class HSwish(nn.Module):
    def forward(self, x):
        out = x * F.relu6(x + 3, inplace=True) / 6
        return out
class ConvBNRelu(nn.Sequential):
    def __init__(self,
                 input_channel,
                 output_channel,
                 kernel,
                 stride,
                 pad,
                 activation="relu",
                 bn=True,
                 group=1,
                 bn_momentum=0.1,
                 *args,
                 **kwargs):

        super(ConvBNRelu, self).__init__()

        assert activation in ["hswish", "relu", None]
        assert stride in [1, 2, 4]

        self.add_module(
            "conv",
            nn.Conv2d(
                input_channel,
                output_channel,
                kernel,
                stride,
                pad,
                groups=group,
                bias=False))
        if bn:
            self.add_module(
                "bn",
                nn.BatchNorm2d(
                    output_channel,
                    momentum=bn_momentum))

        if activation == "relu":
            self.add_module("relu", nn.ReLU6(inplace=True))
        elif activation == "hswish":
            self.add_module("hswish", HSwish())


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = ConvBNRelu(input_channel=3, output_channel=64, kernel=3, stride=1, activation='relu', pad=0,
                                groups=1)
        self.conv2 = ConvBNRelu(input_channel=64, output_channel=64, kernel=5, stride=2, activation='hswish', pad=0,
                                group=1)

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

summary(Model().to('cuda'), (3, 128, 128))

ModulesList

ModulesList是以list的形式保存sub-modules或者网络层，这样就可以先将网络需要的layer构建好保存到一个list，然后通过ModulesList方法添加到网络中。

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

        self.vgg = nn.ModuleList(vgg(vgg_cfg,3))

    def forward(self,x):
        for layer in self.vgg:
            x = layer(x)
print('/输出ModelVGG///')
print(ModelVGG())

读取配置好的网络结构vgg_cfg然后，创建相应的Layer List，使用ModuleList加入到网络中。这样就可以很灵活的创建不同的网络。

使用ModuleList创建类，完成模型构建

import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary


class HSwish(nn.Module):
    def forward(self, x):
        out = x * F.relu6(x + 3, inplace=True) / 6
        return out


class ConvBNRelu(nn.Sequential):
    def __init__(self,
                 input_channel,
                 output_channel,
                 kernel,
                 stride,
                 pad,
                 activation="relu",
                 bn=True,
                 group=1,
                 bn_momentum=0.1,
                 *args,
                 **kwargs):

        super(ConvBNRelu, self).__init__()

        assert activation in ["hswish", "relu", None]
        assert stride in [1, 2, 4]

        self.add_module(
            "conv",
            nn.Conv2d(
                input_channel,
                output_channel,
                kernel,
                stride,
                pad,
                groups=group,
                bias=False))
        if bn:
            self.add_module(
                "bn",
                nn.BatchNorm2d(
                    output_channel,
                    momentum=bn_momentum))

        if activation == "relu":
            self.add_module("relu", nn.ReLU6(inplace=True))
        elif activation == "hswish":
            self.add_module("hswish", HSwish())


class ConvBNReluList(nn.ModuleList):
    def __init__(self,
                 input_channel,
                 output_channel,
                 kernel,
                 stride,
                 pad,
                 activation="relu",
                 bn=True,
                 group=1,
                 bn_momentum=0.1,
                 *args,
                 **kwargs):

        super(ConvBNReluList, self).__init__()

        assert activation in ["hswish", "relu", None]
        assert stride in [1, 2, 4]

        self.add_module(
            "conv",
            nn.Conv2d(
                input_channel,
                output_channel,
                kernel,
                stride,
                pad,
                groups=group,
                bias=False))
        if bn:
            self.add_module(
                "bn",
                nn.BatchNorm2d(
                    output_channel,
                    momentum=bn_momentum))

        if activation == "relu":
            self.add_module("relu", nn.ReLU6(inplace=True))
        elif activation == "hswish":
            self.add_module("hswish", HSwish())


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = ConvBNReluList(input_channel=3, output_channel=64, kernel=3, stride=1, activation='relu', pad=0,
                                groups=1)
        self.conv2 = ConvBNRelu(input_channel=64, output_channel=64, kernel=5, stride=2, activation='hswish', pad=0,
                                group=1)

    def forward(self, x):
        for i, layer in enumerate(self.conv1):
            x = layer(x)
        # x = self.conv1(x)
        x = self.conv2(x)
        return x


summary(Model().to('cuda'), (3, 128, 128))

代码更改class ConvBNRelu(nn.Sequential)所继承的类，为class ConvBNReluList(nn.ModuleList)

在forward函数中，self.conv1使用遍历迭代的方式完成，而self.conv2直接一次性完成。

这里需要注意的是，ModuleList是将Module加入网络中，需要自己手动的遍历进行每一个Module的forward。

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较为方便。对比上面代码：self.conv1和self.conv2就显而易见

总结：

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

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

2、遍历网络结构

可以使用以下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, model in enumerate(m.modules()):
    print(idx, "-", model)

其结果为：

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。

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

/输出网络模型named_modules///
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 - ('conv1', Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1)))
2 - ('conv2', Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)))
3 - ('maxpool1', MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))
4 - ('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()
))
5 - ('features.conv3', Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1)))
6 - ('features.conv4', Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1)))
7 - ('features.relu1', ReLU())

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()
)

print('/输出网络模型named_children///')
for idx, model in enumerate(m.named_children()):
    print(idx, "-", model)

输出为：

/输出网络模型named_children///
0 - ('conv1', Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1)))
1 - ('conv2', Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)))
2 - ('maxpool1', MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))
3 - ('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()
))

3、网络的参数

方法parameters()返回一个包含模型所有参数的迭代器。一般用来当作optimizer的参数。

print('/输出网络参数parameters')
for p in m.parameters():
    print(type(p.data), p.size())
print('/输出网络参数named_parameters//')
for k, v in m.named_parameters():
    print(k, " - ", v.size())

其输出为：

/输出网络参数parameters///
<class 'torch.Tensor'> torch.Size([64, 3, 3, 3])
<class 'torch.Tensor'> torch.Size([64])
<class 'torch.Tensor'> torch.Size([64, 64, 3, 3])
<class 'torch.Tensor'> torch.Size([64])
<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])
/输出网络参数named_parameters///
conv1.weight  -  torch.Size([64, 3, 3, 3])
conv1.bias  -  torch.Size([64])
conv2.weight  -  torch.Size([64, 64, 3, 3])
conv2.bias  -  torch.Size([64])
features.conv3.weight  -  torch.Size([128, 64, 3, 3])
features.conv3.bias  -  torch.Size([128])
features.conv4.weight  -  torch.Size([128, 128, 3, 3])
features.conv4.bias  -  torch.Size([128])

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

named_parameters()返回参数的名称及参数本身，可以按照参数名对一些参数进行处理。

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

4、模型的保存与加载

Pytorch模型加载函数torch.load()_北国觅梦-CSDN博客

PyTorch使用torch.save和torch.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) # 应用到网络结构中

5、加载预训练模型

PyTorch中的torchvision里有很多常用网络的预训练模型，例如：vgg,resnet,googlenet等，可以方便的使用这些预训练模型进行微调。

# PyTorch中的torchvision里有很多常用的模型，可以直接调用：
import torchvision.models as models
 
resnet101 = models.resnet18(pretrained=True)
alexnet = models.alexnet()
squeezenet = models.squeezenet1_0()

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

import torch.utils.model_zoo as model_zoo
model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}
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