pytorch中的pre-train函数模型引用及修改（增减网络层，修改某层参数等）

https://blog.csdn.net/whut_ldz/article/details/78845947

pytorch中的pre-train函数模型引用及修改（增减网络层，修改某层参数等）

2017年12月19日 18:49:37 whut_ldz 阅读数：7421 标签： 深度学习 神经网络 pytorch 预训练 修改 更多

个人分类： python，pytorch，深度学习

一、pytorch中的pre-train模型

卷积神经网络的训练是耗时的，很多场合不可能每次都从随机初始化参数开始训练网络。

pytorch中自带几种常用的深度学习网络预训练模型，如VGG、ResNet等。往往为了加快学习的进度，在训练的初期我们直接加载pre-train模型中预先训练好的参数，model的加载如下所示：

 

1. import torchvision.models as models

2.  

3. #resnet

4. model = models.ResNet(pretrained=True)

5. model = models.resnet18(pretrained=True)

6. model = models.resnet34(pretrained=True)

7. model = models.resnet50(pretrained=True)

8.  

9. #vgg

10. model = models.VGG(pretrained=True)

11. model = models.vgg11(pretrained=True)

12. model = models.vgg16(pretrained=True)

13. model = models.vgg16_bn(pretrained=True)

二、预训练模型的修改

1.参数修改

对于简单的参数修改，这里以resnet预训练模型举例，resnet源代码在Github点击打开链接。

resnet网络最后一层分类层fc是对1000种类型进行划分，对于自己的数据集，如果只有9类，修改的代码如下：

1. # coding=UTF-8

2. import torchvision.models as models

3.  

4. #调用模型

5. model = models.resnet50(pretrained=True)

6. #提取fc层中固定的参数

7. fc_features = model.fc.in_features

8. #修改类别为9

9. model.fc = nn.Linear(fc_features, 9)

2.增减卷积层

前一种方法只适用于简单的参数修改，有的时候我们往往要修改网络中的层次结构，这时只能用参数覆盖的方法，即自己先定义一个类似的网络，再将预训练中的参数提取到自己的网络中来。这里以resnet预训练模型举例。

1. # coding=UTF-8

2. import torchvision.models as models

3. import torch

4. import torch.nn as nn

5. import math

6. import torch.utils.model_zoo as model_zoo

7.  

8. class CNN(nn.Module):

9.  

10. def __init__(self, block, layers, num_classes=9):

11. self.inplanes = 64

12. super(ResNet, self).__init__()

13. self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,

14. bias=False)

15. self.bn1 = nn.BatchNorm2d(64)

16. self.relu = nn.ReLU(inplace=True)

17. self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

18. self.layer1 = self._make_layer(block, 64, layers[0])

19. self.layer2 = self._make_layer(block, 128, layers[1], stride=2)

20. self.layer3 = self._make_layer(block, 256, layers[2], stride=2)

21. self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

22. self.avgpool = nn.AvgPool2d(7, stride=1)

23. #新增一个反卷积层

24. self.convtranspose1 = nn.ConvTranspose2d(2048, 2048, kernel_size=3, stride=1, padding=1, output_padding=0, groups=1, bias=False, dilation=1)

25. #新增一个最大池化层

26. self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)

27. #去掉原来的fc层，新增一个fclass层

28. self.fclass = nn.Linear(2048, num_classes)

29.  

30. for m in self.modules():

31. if isinstance(m, nn.Conv2d):

32. n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels

33. m.weight.data.normal_(0, math.sqrt(2. / n))

34. elif isinstance(m, nn.BatchNorm2d):

35. m.weight.data.fill_(1)

36. m.bias.data.zero_()

37.  

38. def _make_layer(self, block, planes, blocks, stride=1):

39. downsample = None

40. if stride != 1 or self.inplanes != planes * block.expansion:

41. downsample = nn.Sequential(

42. nn.Conv2d(self.inplanes, planes * block.expansion,

43. kernel_size=1, stride=stride, bias=False),

44. nn.BatchNorm2d(planes * block.expansion),

45. )

46.  

47. layers = []

48. layers.append(block(self.inplanes, planes, stride, downsample))

49. self.inplanes = planes * block.expansion

50. for i in range(1, blocks):

51. layers.append(block(self.inplanes, planes))

52.  

53. return nn.Sequential(*layers)

54.  

55. def forward(self, x):

56. x = self.conv1(x)

57. x = self.bn1(x)

58. x = self.relu(x)

59. x = self.maxpool(x)

60.  

61. x = self.layer1(x)

62. x = self.layer2(x)

63. x = self.layer3(x)

64. x = self.layer4(x)

65.  

66. x = self.avgpool(x)

67. #新加层的forward

68. x = x.view(x.size(0), -1)

69. x = self.convtranspose1(x)

70. x = self.maxpool2(x)

71. x = x.view(x.size(0), -1)

72. x = self.fclass(x)

73.  

74. return x

75.  

76. #加载model

77. resnet50 = models.resnet50(pretrained=True)

78. cnn = CNN(Bottleneck, [3, 4, 6, 3])

79. #读取参数

80. pretrained_dict = resnet50.state_dict()

81. model_dict = cnn.state_dict()

82. # 将pretrained_dict里不属于model_dict的键剔除掉

83. pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}

84. # 更新现有的model_dict

85. model_dict.update(pretrained_dict)

86. # 加载我们真正需要的state_dict

87. cnn.load_state_dict(model_dict)

88. # print(resnet50)

89. print(cnn)

以上就是相关的内容，本人刚入门的小白一枚，请轻喷～