Pytorch – 模型微调时删除原有模型中的某一层的方法

本文以去除Pytorch预置的ResNet18网络中最后一层全连接分类层为例，说明模型微调时如何去除模型中某一层的方法。

我们想要在模型中去掉某一层实际上就等效于在该层不进行任何操作，直接将上一层的值直接返回即可，下面提供了3种方法进行选择。

1 使用自定义nn.Module替换指定层

在本方法中，自定义了一个Identity类继承自nn.Module，并在forward方法中不进行如何操作直接返回原值，达到去除某一层的效果。

# -*- coding: utf-8 -*-

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

class Identity(nn.Module):
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, x):
        return x

if __name__ == '__main__':
    resnet18_modify = torchvision.models.resnet18(pretrained=True)
    resnet18_modify.fc = Identity()

    tensor = torch.rand([1,3,244,244])

    out = resnet18_modify(tensor)

    print(out.shape)

2 将需删除的层指定为空的nn.Sequential

这种方法更加简单，直接将想要删除的层设置为空的nn.Sequential，同样是不进行任何操作返回原值。

# -*- coding: utf-8 -*-

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

if __name__ == '__main__':
    resnet18_modify = torchvision.models.resnet18(pretrained=True)
    resnet18_modify.fc = nn.Sequential()

    tensor = torch.rand([1,3,244,244])

    out = resnet18_modify(tensor)

    print(out.shape)

3 先将已有模型中的子模块放到list中，然后在list中去掉某一层，然后再重组网络

# -*- coding: utf-8 -*-

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

if __name__ == '__main__':
    resnet18_modify = torchvision.models.resnet18(pretrained=True)
    print(resnet18_modify)

    modules = list(resnet18_modify.children())[:-1]
    resnet18_modify = nn.Sequential(*modules)

    print(resnet18_modify)

    tensor = torch.rand([1,3,244,244])

    out = resnet18_modify(tensor)

    print(out.shape)