一、pytorch中的pre-train模型

卷积神经网络的训练是耗时的，很多场合不可能每次都从随机初始化参数开始训练网络。
pytorch中自带几种常用的深度学习网络预训练模型，如VGG、ResNet等。往往为了加快学习的进度，在训练的初期我们直接加载pre-train模型中预先训练好的参数，model的加载如下所示：

  1. import torchvision.models as models
  3. #resnet
  4. model = models.ResNet(pretrained=True)
  6. #vgg
  7. model = models.VGG(pretrained=True)

二、预训练模型的修改

1. 参数修改

对于简单的参数修改，这里以resnet预训练模型举例，resnet源代码在Github点击打开链接。
resnet网络最后一层分类层fc是对1000种类型进行划分，对于自己的数据集，如果只有9类，修改的代码如下：

  1. # coding=UTF-8
  2. import torchvision.models as models
  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
  8. class CNN(nn.Module):
  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, bias=False)
  14.         self.bn1 = nn.BatchNorm2d(64)
  15.         self.relu = nn.ReLU(inplace=True)
  16.         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
  17.         self.layer1 = self._make_layer(block, 64, layers[0])
  18.         self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
  19.         self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
  20.         self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
  21.         self.avgpool = nn.AvgPool2d(7, stride=1)
  22.         # 新增一个反卷积层
  23.         self.convtranspose1 = nn.ConvTranspose2d(2048, 2048, kernel_size=3, stride=1, padding=1, output_padding=0, groups=1, bias=False, dilation=1)
  24.         # 新增一个最大池化层
  25.         self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
  26.         # 去掉原来的fc层，新增一个fclass层
  27.         self.fclass = nn.Linear(2048, num_classes)
  29.         for m in self.modules():
  30.             if isinstance(m, nn.Conv2d):
  31.                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
  32.                 m.weight.data.normal_(0, math.sqrt(2. / n))
  33.             elif isinstance(m, nn.BatchNorm2d):
  34.                 m.weight.data.fill_(1)
  35.                 m.bias.data.zero_()
  37.     def _make_layer(self, block, planes, blocks, stride=1):
  38.         downsample = None
  39.         if stride != 1 or self.inplanes != planes * block.expansion:
  40.             downsample = nn.Sequential(
  41.                 nn.Conv2d(self.inplanes, planes * block.expansion,
  42.                           kernel_size=1, stride=stride, bias=False),
  43.                 nn.BatchNorm2d(planes * block.expansion),
  44.             )
  46.         layers = []
  47.         layers.append(block(self.inplanes, planes, stride, downsample))
  48.         self.inplanes = planes * block.expansion
  49.         for i in range(1, blocks):
  50.             layers.append(block(self.inplanes, planes))
  52.         return nn.Sequential(*layers)
  54.     def forward(self, x):
  55.         x = self.conv1(x)
  56.         x = self.bn1(x)
  57.         x = self.relu(x)
  58.         x = self.maxpool(x)
  60.         x = self.layer1(x)
  61.         x = self.layer2(x)
  62.         x = self.layer3(x)
  63.         x = self.layer4(x)
  65.         x = self.avgpool(x)
  66.         # 新加层的forward
  67.         x = x.view(x.size(0), -1)
  68.         x = self.convtranspose1(x)
  69.         x = self.maxpool2(x)
  70.         x = x.view(x.size(0), -1)
  71.         x = self.fclass(x)
  73.         return x
  75. # 加载model
  76. resnet50 = models.resnet50(pretrained=True)  # 官方的pre模型
  77. cnn = CNN(Bottleneck, [3, 4, 6, 3])  # 自定义的模型
  78. # 读取参数
  79. pretrained_dict = resnet50.state_dict()
  80. model_dict = cnn.state_dict()
  81. # 将pretrained_dict里不属于model_dict的键剔除掉
  82. pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
  83. # 更新现有的model_dict
  84. model_dict.update(pretrained_dict)
  85. # 加载我们真正需要的state_dict
  86. cnn.load_state_dict(model_dict)
  87. # print(resnet50)
  88. print(cnn)

三、在多GPU集群上加载单GPU的model

  1. ######### multi-gpu load sigle-gpu need `mudule` ########### 
  2. pretrained_dict = torch.load("model_ir_se50.pth")  # pre-trained model
  3. self.model_dict = self.model.state_dict()  #get the name:value
  4. param={}
  5. for k, v in pretrained_dict.items():
  6.     if k[7:] !="module" :
  7.         param["module."+k] = pretrained_dict[k]
  8.         pretrained_dict = {k: v for k, v in param.items() if k in self.model_dict}
  9.         self.model_dict.update(pretrained_dict)
  10.         self.model.load_state_dict(self.model_dict)

实际上就是字典的操作，那么字典的操作，哪些层不要，打印出层的名字就可以了，例如
模型参数的某些层的权重不要，那么重构一个字典参数就可以了，for k in torch.load(“**.pth’).keys(): 打印出来按照名字删除 键值对，
所以字典的pop删除操作也是可以的， 多卡训练参数多了一个module

  1. ######### sigle-gpu load multi-gpu`s model need to remove `mudule.` ########### 
  2. pretrain = torch.load("/home/imagenet.pth")
  3. new_state_dict = {}  # OrderedDict()
  4. for k, v in pretrain.items():
  5.     if "classifier" in k:  #最后分类层的参数是classeifer, 不需要这个模型参数
  6.         continue
  7.     new_state_dict[k[7:]] = v  #remove `module.`  #模型k 有module 不要
  8. model.load_state_dict(new_state_dict, strict=False)  #strict =False ,模型参数和模型不一致可以加载