torch 迁移学习固定指定部分权重

最新推荐文章于 2024-07-13 14:57:04 发布

蓝羽飞鸟

最新推荐文章于 2024-07-13 14:57:04 发布

阅读量263

点赞数 1

文章标签：迁移学习深度学习人工智能

本文链接：https://blog.csdn.net/level_code/article/details/130553372

版权

文章讲述了在迁移学习中如何控制模型的训练过程，特别是如何通过设置`requires_grad`属性来决定模型的哪些部分参与训练。以ResNet18为例，展示了如何遍历模型的子模块，将前6层及第6层的第一个子模块的权重固定，使其不参与训练。这通常在使用预训练模型并微调特定部分时进行。

摘要由CSDN通过智能技术生成

在迁移学习中，如果想只训练其中一部分，那么需要固定指定部分的权重。
模型中，通过指定参数的requires_grad来确定是否需要训练。
指定requires_grad为True时，参与训练，否则不参与训练。

现在需要知道模型中都有哪些module, 挑选出参与 / 不参与训练的部分，设置requires_grad即可。

通过model.children()获得child模块，再通过child.parameters()获得参数param,
通过param.requires_grad来设置。

以resnet18为例，先看看都有哪些模块。

child_counter = 0
for child in model.children():
    print(" child", child_counter)
    print(child)
    child_counter += 1

child 0
Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
child 1
BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True)
child 2
ReLU (inplace)
child 3
MaxPool2d (size=(3, 3), stride=(2, 2), padding=(1, 1), dilation=(1, 1))
child 4
Sequential (
  (0): BasicBlock (
    (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True)
    (relu): ReLU (inplace)
    (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True)
  )
  (1): BasicBlock (
    (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True)
    (relu): ReLU (inplace)
    (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True)
  )
)
...

如果有预训练好的模型，先load, 再设置requires_grad.
比如说想把child<6处的模块固定住，不参与训练。
同时还想把child=6处模块的第一个子模块固定住。

from torchvision import datasets, models, transforms

#有checkpoint的要先load
model = models.resnet18(pretrained = False)
checkpoint = torch.load(MODEL_PATH)
model.load_state_dict(checkpoint)

child_counter = 0
for child in model.children():
    if child_counter < 6:
        for param in child.parameters():
            param.requires_grad = False
    elif child_counter == 6:
        children_of_child_counter = 0
        for children_of_child in child.children():
            if children_of_child_counter < 1:
                for param in children_of_child.parameters():
                    param.requires_grad = False
            children_of_child_counter += 1
    child_counter += 1