Pytorch学习笔记（I）——预训练模型（二）：修改网络结构（ResNet50及以上/VGG16）_resnet50的后两层

（pytorch1.0）最近在研究pytorch如何修改与训练模型的网络结构，然后发现了两种版本，一种是细调版，一种是快速版
　　经过一番钻研后发现细调版适合对网络模型进行大幅度的改动（如在原有的结构上穿插着增减层），而快速版适合直接对网络末端的层进行增减。
　　虽然快速版简单易懂，但是还是要对细调版有所了解才能比较，万一以后用的上呢。因此，我就好好研究了一番细调版，结果发现网上的代码或者博客基本都是相互搬运的，代码中的错误一模一样，对于我这种小白来说特别不友好。于是，我就在前人的基础上查缺补漏，重新整理了一下。
　　关于如何加载和使用，请查看前一篇博客

以下示例增减的层结构不一定合适，仅供参考

一、细调版

话不多说，直接上代码，这里以resnet50为例
这一步必须参考原来的网络结构，从而定义一个类似的网络

import torchvision.models as models
import torch
import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo

#Bottleneck是一个class 里面定义了使用1*1的卷积核进行降维跟升维的一个残差块，可以在github resnet pytorch上查看
class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

#不做修改的层不能乱取名字，否则预训练的权重参数无法传入
class CNN(nn.Module):

    def __init__(self, block, layers, num_classes=9):
        self.inplanes = 64
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1,1))
        # 新增一个反卷积层
        self.convtranspose1 = nn.ConvTranspose2d(2048, 2048, kernel_size=3, stride=1, padding=1, output_padding=0,
                                                 groups=1, bias=False, dilation=1)
        # 新增一个最大池化层
        self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        # 去掉原来的fc层，新增一个fclass层
        self.fclass = nn.Linear(2048, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)
	#这一步用以设置前向传播的顺序，可以自行调整，前提是合理
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        # 新加层的forward
        x = x.view(x.size(0), -1)
        x = self.convtranspose1(x)
        x = self.maxpool2(x)
        x = x.view(x.size(0), -1)
        x = self.fclass(x)

        return x


# 加载model
resnet50 = models.resnet50(pretrained=True)
#3 4 6 3 分别表示layer1 2 3 4 中Bottleneck模块的数量。res101则为3 4 23 3 
cnn = CNN(Bottleneck, [3, 4, 6, 3])
# 读取参数
pretrained_dict = resnet50.state_dict()
model_dict = cnn.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
cnn.load_state_dict(model_dict)
# print(resnet50)
print(cnn)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

接下来我们来比对一下前后的变化。
1、 原来的resnet50最后两层信息如下

(avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
(fc): Linear(in_features=2048, out_features=1000, bias=True)
1
2

2、 新的最后几层层信息如下

(avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
(convtranspose1): ConvTranspose2d(2048, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(maxpool2): MaxPool2d(kernel_size=3, stride=1, padding=1, dilation=1, ceil_mode=False)
(fclass): Linear(in_features=2048, out_features=9, bias=True)
1
2
3
4

可以看出在最后一层fc被替换了。
如果想在中间进行增减，改变前向传播的顺序就好了

二、快速版

这里以vgg16为例

import torchvision.models as models
import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self, model):
        super(Net, self).__init__()
        # -2表示去掉model的后两层
        self.vgg_layer = nn.Sequential(*list(model.children())[:-2])
        self.transion_layer = nn.ConvTranspose2d(2048, 2048, kernel_size=14, stride=3)
        self.pool_layer = nn.MaxPool2d(32)
        self.Linear_layer = nn.Linear(2048, 8)

    def forward(self, x):
        x = self.vgg_layer(x)
        x = self.transion_layer(x)
        x = self.pool_layer(x)
        #将一个多行的Tensor,拼接成一行,-1指在不告诉函数有多少列
        x = x.view(x.size(0), -1)
        x = self.Linear_layer(x)
        return x

vgg = models.vgg16(pretrained=True)
model = Net(vgg)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

1、 原来的vgg16特征提取之后有一个pooling层和分类器

(avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
(classifier): Sequential(
  (0): Linear(in_features=25088, out_features=4096, bias=True)
  (1): ReLU(inplace)
  (2): Dropout(p=0.5)
  (3): Linear(in_features=4096, out_features=4096, bias=True)
  (4): ReLU(inplace)
  (5): Dropout(p=0.5)
  (6): Linear(in_features=4096, out_features=1000, bias=True)
1
2
3
4
5
6
7
8
9

2、 新的最后几层层信息如下

(transion_layer): ConvTranspose2d(2048, 2048, kernel_size=(14, 14), stride=(3, 3))
(pool_layer): MaxPool2d(kernel_size=32, stride=32, padding=0, dilation=1, ceil_mode=False)
(Linear_layer): Linear(in_features=2048, out_features=8, bias=True)
1
2
3

我用附录的代码查看后，发现权值已经载入到新的模型了，不用像细调版那样。

附录、查看权值

#查看vgg16
for child in vgg.modules():
    print(child)
    for param in child.parameters():
        print(param)
        
 #查看新的Net
 for child in modules.modules():
    print(child)
    for param in child.parameters():
        print(param)
1
2
3
4
5
6
7
8
9
10
11