如何在 PyTorch 中冻结模型权重以进行迁移学习：分步教程

news2024/11/19 22:53:32

一、说明

迁移学习是一种机器学习技术，其中预先训练的模型适用于新的但类似的问题。迁移学习的关键步骤之一是能够冻结预训练模型的层，以便在训练期间仅更新网络的某些部分。当您想要保留预训练模型已经学习的特征时，冻结至关重要。在本教程中，我们将使用一个简单的示例来演示在 PyTorch 中冻结权重以进行迁移学习的过程。

二、先决条件

如果您没有安装 torch 和 torchvision 库，我们可以在终端中执行以下操作：

pip install torch torchvision

三、导入库

让我们从 Python 代码开始。首先，我们导入本教程的库：

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

四、加载预训练模型

我们将在此示例中使用预训练的 ResNet-18 模型：

# Load the pre-trained model
resnet18 = models.resnet18(pretrained=True)

五、冻结层

要冻结图层，我们将requires_grad属性设置为False。这可以防止 PyTorch 在反向传播期间计算这些层的梯度。

# Freeze all layers
for param in resnet18.parameters():
    param.requires_grad = False

六、解冻一些层

通常，为了获得最佳结果，我们会对网络中的后续层进行一些微调。我们可以这样做：

# Unfreeze last layer
for param in resnet18.fc.parameters():
    param.requires_grad = True

七、修改网络架构

我们将替换最后一个全连接层，以使模型适应具有不同数量的输出类（假设有 10 个类）的新问题。此外，这使我们能够将这个预训练网络用于分类以外的其他应用，例如分割。对于分割，我们用卷积层替换最后一层。对于此示例，我们继续执行包含 10 个类别的分类任务。

# Replace last layer
num_ftrs = resnet18.fc.in_features
resnet18.fc = nn.Linear(num_ftrs, 10)

八、训练修改后的模型

让我们定义一个简单的训练循环。出于演示目的，我们将使用随机数据。

# Create random data
inputs = torch.randn(5, 3, 224, 224)
labels = torch.randint(0, 10, (5,))

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(resnet18.fc.parameters(), lr=0.001, momentum=0.9)

# Training loop
for epoch in range(5):
    optimizer.zero_grad()
    outputs = resnet18(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    print(f'Epoch {epoch+1}/5, Loss: {loss.item()}')

在此示例中，训练期间仅更新最后一层的权重。