迁移学习

迁移学习

第一种做法

把模型的参数加载到自己的网络中
Net部分

import math
import torch.nn as nn
from torchvision.models import ResNet
from torchvision import models
import torch


def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = 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:
            identity = self.downsample(x)

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

        return out


class CNN(nn.Module):
    def __init__(self, block, layers, num_classes=120):

        super(CNN, self).__init__()  
        self.inplanes = 64
        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.AvgPool2d(7, stride=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(4 * 4 * 2048,
                                num_classes)  # 原来的fc层：self.fc = nn.Linear(512 * block.expansion, 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)
        x = self.convtranspose1(x)
        x = self.maxpool2(x)
        x = x.reshape(-1, 4 * 4 * 2048)
        x = self.fclass(x)

        return x


if __name__ == '__main__':
	'''后面接了全连接层，因此要改变图像的输入，则需要改变全连接层'''
    x = torch.randn(1, 3, 299, 299)
    net = CNN(Bottleneck, [3, 4, 6, 3])
    print(net(x).shape)

from torchvision import models
import torch

model = Net()  # 这是自定义网络
# 当权重信息在本地
resnet50 = models.resnet50(pretrained=False)
# 本地权重路径
pretrained_file = "model/resnet50.pth"
# 加载参数
resnet50.load_state_dict(torch.load(pretrained_file))
# 本地无权重信息时
# resnet50 = models.resnet50(pretrained=True)
# resnet50.load_state_dict(torch.load(resnet50))
# 加载
pretrained_dict = resnet50.state_dict()  # 记录预训练模型的参数 
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.updata(pretrained_dict)
# 加载我们需要的state_dict
model.load_state_dict(model_dict)
for param in model.parameters():
    param.requires_grad = True  # 让参数可训练

第二种做法

直接在训练最后一层

from torchvision import models
import torch.nn as nn
model = models.inception_v3(pretrained=True)
for param in model.parameters():
    param.requires_grad = False  # 让参数不可被训练
n_classes = 120  # 类别个数
n_inputs = model.fc.in_features  # 源网络全连接层输出
model.fc = nn.Sequential(
	nn.Linear(n_inputs, 1024),
	nn.Relu(),
    nn.Dropout(.4),
    nn.Linear(1024, n_classes),
    nn.LogSoftmax(dim=1)
)