pytorch模型参数迁移(三种方法)

1.利用resnet18做迁移学习

1.先下载训练好的resnet18-5c106cde.pth,并存在指定的位置上

import torchvision.models as models
 
# pretrained=True就可以使用预训练的模型
resnet18 = models.resnet18(pretrained=True)
torch.save(model.state_dict(), 'model/resnet18-5c106cde.pth')

2.加载resnet18的模型和运行

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


class BasicBlock(nn.Module):# 针对于renet18至34层的残差结构
    expansion = 1
    def __init__(self, in_channel, out_channel, stride=1, downsample=None,**kwargs):
        super(BasicBlock, self).__init__()

        self.conv1 = nn.Conv2d(in_channels= in_channel, out_channels=out_channel, kernel_size=3,
                                stride =stride,padding=1,bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(in_channels= out_channel, out_channels=out_channel, kernel_size=3,
                               stride =1,padding=1,bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.downsample = downsample

    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)

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

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

        return out

class Bottleneck(nn.Module):
    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
    # This variant is also known as ResNet V1.5 and improves accuracy according to
    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.

    expansion = 4

    def __init__(self, in_channel, out_channel, stride=1, downsample=None,
                 groups= 1, width_per_group= 64):
        super(Bottleneck, self).__init__()

        self.conv1 = nn.Conv2d(in_channels= in_channel, out_channels=out_channel, kernel_size=1,
                                stride =stride,padding=1,bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.conv2 = nn.Conv2d(in_channels= out_channel, out_channels=out_channel, kernel_size=3,
                                stride =stride,padding=1,bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.conv3 = nn.Conv2d(in_channels= out_channel, out_channels=out_channel*self.expansion, kernel_size=1,
                                stride =1,bias=False)
        self.bn3 = nn.BatchNorm2d(out_channel * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample


    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 ResNet(nn.Module):

    def __init__(self, num_outputs=None,  #输出的分类数
                 backbone=None,
                 pretrained=False,
                 curriculum_steps=None,
                 extra_outputs=0,
                 share_top_y=True,
                 pred_category=False,
                 block=BasicBlock, block_num=[2,2,2,2],include_top=True, groups=1,
                 width_per_group=64):
        # blocks_num:残差结构中每个block存在多少个layer层
        super(ResNet, self).__init__()

        self.include_top = include_top
        self.in_channel = 64  # 输入图片经过第一层卷积的通道数
        self.groups = groups
        self.width_per_group = width_per_group
        self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(self.in_channel)
        self.relu = nn.ReLU(inplace=True)

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

        self.layer1 = self._make_layer(block, 64, block_num[0])
        self.layer2 = self._make_layer(block, 128, block_num[1], stride=2)
        self.layer3 = self._make_layer(block, 256, block_num[2], stride=2)
        self.layer4 = self._make_layer(block, 512, block_num[3], stride=2)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # output size =(1,1)
        self.fc = nn.Linear(512 * block.expansion, 1000)

        image_size = 12
        patch_size = 3  # 后期尝试改为2
        dim = 128
        depth = 2
        num_classes = 35
        expansion_factor = 4
        num_patches = (image_size // patch_size) ** 2


        self.curriculum_steps = [0, 0, 0, 0] if curriculum_steps is None else curriculum_steps
        self.share_top_y = share_top_y
        self.extra_outputs = extra_outputs
        self.pred_category = pred_category
        self.sigmoid = nn.Sigmoid()

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

        layers = []
        layers.append(block(self.in_channel,
                            channel,
                            downsample=downsample,
                            stride=stride,
                            groups=self.groups,
                            width_per_group=self.width_per_group))
        self.in_channel = channel * block.expansion
        for _ in range(1, block_num):
            layers.append(block(self.in_channel,
                                channel,
                                groups=self.groups,
                                width_per_group=self.width_per_group))

        return nn.Sequential(*layers)

    def forward(self, x, epoch=None, **kwargs):

        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)  # torch.Size[B 128 12 20]

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x

if __name__ == "__main__":
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device = 'cpu'
    print("-----device:{}".format(device))
    print("-----Pytorch version:{}".format(torch.__version__))

    input_tensor = torch.zeros(1, 3, 100, 100)
    print('input_tensor:', input_tensor.shape)



    pretrained_file = "./model_resnet18.pt"
    model = ResNet()
    model.load_state_dict(torch.load(pretrained_file))
    model.eval()
    out = model(input_tensor)
    print("out:", out.shape, out[0, 0:10])

运行结果如下:

-----device:cpu
-----Pytorch version:1.5.0
input_tensor: torch.Size([1, 3, 100, 100])
out: torch.Size([1, 1000]) tensor([ 0.4010,  0.8436,  0.3071,  0.0627,  0.4446,  0.8470,  0.1882,  0.7012,
         0.2988, -0.7574], grad_fn=<SliceBackward>)

3.修改resnet18的网络架构后,如何加载原来已经训练好的模型参数。
例如:

#将114行的代码修改成
self.layer44 = self._make_layer(block, 512, block_num[3], stride=2)
#将166行的代码修改成
x = self.layer44(x)

直接加载模型,运行结果:

RuntimeError: Error(s) in loading state_dict for ResNet:
	Missing key(s) in state_dict: "layer44.0.conv1.weight", "layer44.0.bn1.weight", "layer44.0.bn1.bias", "layer44.0.bn1.running_mean", "layer44.0.bn1.running_var", "layer44.0.conv2.weight", "layer44.0.bn2.weight", "layer44.0.bn2.bias", "layer44.0.bn2.running_mean", "layer44.0.bn2.running_var", "layer44.0.downsample.0.weight", "layer44.0.downsample.1.weight", "layer44.0.downsample.1.bias", "layer44.0.downsample.1.running_mean", "layer44.0.downsample.1.running_var", "layer44.1.conv1.weight", "layer44.1.bn1.weight", "layer44.1.bn1.bias", "layer44.1.bn1.running_mean", "layer44.1.bn1.running_var", "layer44.1.conv2.weight", "layer44.1.bn2.weight", "layer44.1.bn2.bias", "layer44.1.bn2.running_mean", "layer44.1.bn2.running_var". 
	Unexpected key(s) in state_dict: "layer4.0.conv1.weight", "layer4.0.bn1.weight", "layer4.0.bn1.bias", "layer4.0.bn1.running_mean", "layer4.0.bn1.running_var", "layer4.0.bn1.num_batches_tracked", "layer4.0.conv2.weight", "layer4.0.bn2.weight", "layer4.0.bn2.bias", "layer4.0.bn2.running_mean", "layer4.0.bn2.running_var", "layer4.0.bn2.num_batches_tracked", "layer4.0.downsample.0.weight", "layer4.0.downsample.1.weight", "layer4.0.downsample.1.bias", "layer4.0.downsample.1.running_mean", "layer4.0.downsample.1.running_var", "layer4.0.downsample.1.num_batches_tracked", "layer4.1.conv1.weight", "layer4.1.bn1.weight", "layer4.1.bn1.bias", "layer4.1.bn1.running_mean", "layer4.1.bn1.running_var", "layer4.1.bn1.num_batches_tracked", "layer4.1.conv2.weight", "layer4.1.bn2.weight", "layer4.1.bn2.bias", "layer4.1.bn2.running_mean", "layer4.1.bn2.running_var", "layer4.1.bn2.num_batches_tracked".

方法一:将原来预训练好的模型参数迁移到新的resnet18网络架构中,只有迁移两者相同的模型参数,不同的参数还是随机初始化。

def transfer_model(pretrained_file, model):

    pretrained_dict = torch.load(pretrained_file)  # get pretrained dict
    model_dict = model.state_dict()  # get model dict
    # 在合并前(update),需要去除pretrained_dict一些不需要的参数
    pretrained_dict = transfer_state_dict(pretrained_dict, model_dict)
    model_dict.update(pretrained_dict)  # 更新(合并)模型的参数
    model.load_state_dict(model_dict)
    return model


def transfer_state_dict(pretrained_dict, model_dict):
    # state_dict2 = {k: v for k, v in save_model.items() if k in model_dict.keys()}
    state_dict = {}
    for k, v in pretrained_dict.items():
        if k in model_dict.keys():
            # state_dict.setdefault(k, v)
            state_dict[k] = v
        else:
            print("Missing key(s) in state_dict :{}".format(k))
    return state_dict


if __name__ == "__main__":
    input_tensor = torch.zeros(1, 3, 100, 100)
    print('input_tensor:', input_tensor.shape)
    pretrained_file = "./model_resnet18.pt"
    # model = resnet18()
    # model.load_state_dict(torch.load(pretrained_file))
    # model.eval()
    # out = model(input_tensor)
    # print("out:", out.shape, out[0, 0:10])

    model1 = ResNet()
    model1 = transfer_model(pretrained_file, model1)
    out1 = model1(input_tensor)
    print("out1:", out1.shape, out1[0, 0:10])

方法二:修改网络名称并迁移学习
由于我们将官方的resnet18的self.layer4改为了:self.layer44 ,我们仅仅修改了一个网络名称而已,就导致模型参数加载出错。那么,我们如何将预训练好的模型修改成符合新网络架构?

def string_rename(old_string, new_string, start, end):
    new_string = old_string[:start] + new_string + old_string[end:]
    return new_string


def modify_model(pretrained_file, model, old_prefix, new_prefix):
    '''
    :param pretrained_file:
    :param model:
    :param old_prefix:
    :param new_prefix:
    :return:
    '''
    pretrained_dict = torch.load(pretrained_file)
    model_dict = model.state_dict()
    state_dict = modify_state_dict(pretrained_dict, model_dict, old_prefix, new_prefix)
    model.load_state_dict(state_dict)
    return model


def modify_state_dict(pretrained_dict, model_dict, old_prefix, new_prefix):
    '''
    修改model dict
    :param pretrained_dict:
    :param model_dict:
    :param old_prefix:
    :param new_prefix:
    :return:
    '''
    state_dict = {}
    for k, v in pretrained_dict.items():
        if k in model_dict.keys():
            # state_dict.setdefault(k, v)
            state_dict[k] = v
        else:
            for o, n in zip(old_prefix, new_prefix):
                prefix = k[:len(o)]
                if prefix == o:
                    kk = string_rename(old_string=k, new_string=n, start=0, end=len(o))
                    print("rename layer modules:{}-->{}".format(k, kk))
                    state_dict[kk] = v
    return state_dict


if __name__ == "__main__":
    input_tensor = torch.zeros(1, 3, 100, 100)
    print('input_tensor:', input_tensor.shape)
    pretrained_file = "./model_resnet18.pt"
    new_file = "./model_resnet18_1.pt"
    model = ResNet()
    new_model = modify_model(pretrained_file, model, old_prefix=["layer4"], new_prefix=["layer44"])
    torch.save(new_model.state_dict(), new_file)

    model2 = ResNet()
    model2.load_state_dict(torch.load(new_file))
    model2.eval()
    out2 = model2(input_tensor)
    print("out2:", out2.shape, out2[0, 0:10])

方法三:去除原模型的某些模块
在对resnet18的实际应用,一般需要子模块"fc"和"avgpool"来应对实际问题。下面在不修改原模型的基础上,,通过"resnet18.named_children()"和"resnet18.children()“的方法去除子模块"fc"和"avgpool”。

import torch
import torchvision.models as models
from collections import OrderedDict
 
if __name__=="__main__":
  resnet18 = models.resnet18(False)
  print("resnet18",resnet18)
 
  # use named_children()
  resnet18_v1 = OrderedDict(resnet18.named_children())
  # remove avgpool,fc
  resnet18_v1.pop("avgpool")
  resnet18_v1.pop("fc")
  resnet18_v1 = torch.nn.Sequential(resnet18_v1)
  print("resnet18_v1",resnet18_v1)
  # use children
  resnet18_v2 = torch.nn.Sequential(*list(resnet18.children())[:-2])
  print(resnet18_v2,resnet18_v2)
  • 4
    点赞
  • 25
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
风格迁移是一种将一张图像的内容与另一张图像的风格合成在一起的技术。在PyTorch中,可以使用预训练的神经网络模型来实现风格迁移。以下是一个简单的实现步骤: 1. 加载预训练的模型 PyTorch中有许多预训练的模型,包括VGG等,可以用于风格迁移。可以使用torchvision包来加载VGG模型: ```python import torch import torchvision.models as models model = models.vgg19(pretrained=True).features ``` 2. 定义损失函数 风格迁移的损失函数包括内容损失和风格损失。内容损失用于保留原始图像的内容,风格损失用于迁移另一张图像的风格。可以使用MSELoss函数来计算损失: ```python mse_loss = torch.nn.MSELoss() ``` 3. 定义优化器 可以使用Adam优化器来优化图像: ```python optimizer = torch.optim.Adam([target], lr=0.01) ``` 其中,target是生成的图像。 4. 迭代优化 在每次迭代中,使用VGG模型计算目标图像的内容和风格,并计算损失。然后,使用优化器来更新目标图像,以最小化损失。 ```python for i in range(num_iterations): optimizer.zero_grad() target_features = model(target) content_loss = mse_loss(target_features[layer], content_features[layer]) style_loss = mse_loss(gram_matrix(target_features[layer]), gram_matrix(style_features[layer])) loss = alpha * content_loss + beta * style_loss loss.backward() optimizer.step() ``` 其中,layer是VGG模型中的某一层,content_features和style_features分别是原始图像和风格图像在该层中的特征,alpha和beta是内容损失和风格损失的权重。 5. 输出结果 最终生成的图像就是目标图像。可以使用PIL库将其保存到本地: ```python from PIL import Image output_image = target.detach().squeeze(0).permute(1, 2, 0).numpy() output_image = (output_image * 255).clip(0, 255).astype("uint8") output_image = Image.fromarray(output_image) output_image.save("output.jpg") ``` 这就是用PyTorch实现风格迁移的基本步骤。当然,还有很多细节和参数需要调整,具体可以根据实际情况进行调整。
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值