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

一、说明

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

二、先决条件

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

aba 复制代码
pip install torch torchvision 

三、导入库

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

aba 复制代码
import torch
import torch.nn as nn
import torchvision.models as models

四、加载预训练模型

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

aba 复制代码
# Load the pre-trained model
resnet18 = models.resnet18(pretrained=True)

五、冻结层

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

aba 复制代码
# Freeze all layers
for param in resnet18.parameters():
    param.requires_grad = False

六、解冻一些层

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

aba 复制代码
# Unfreeze last layer
for param in resnet18.fc.parameters():
    param.requires_grad = True

七、修改网络架构

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

aba 复制代码
# Replace last layer
num_ftrs = resnet18.fc.in_features
resnet18.fc = nn.Linear(num_ftrs, 10)

八、训练修改后的模型

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

aba 复制代码
# 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()}')  

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

九、结论

在 PyTorch 中冻结层非常简单明了。通过将该requires_grad属性设置为False,您可以防止在训练期间更新特定层,从而使您能够有效地利用预训练模型的强大功能。

了解如何在 PyTorch 中冻结和解冻层对于有效的迁移学习至关重要,因为它允许您利用预训练的模型来执行类似但不同的任务。通过这种简单而强大的技术,您可以在训练深度神经网络时节省时间和计算资源。

参考资料:请访问此处、GithubLinkedIn礼萨·卡兰塔尔

相关推荐
董厂长3 小时前
langchain :记忆组件混淆概念澄清 & 创建Conversational ReAct后显示指定 记忆组件
人工智能·深度学习·langchain·llm
G皮T6 小时前
【人工智能】ChatGPT、DeepSeek-R1、DeepSeek-V3 辨析
人工智能·chatgpt·llm·大语言模型·deepseek·deepseek-v3·deepseek-r1
九年义务漏网鲨鱼6 小时前
【大模型学习 | MINIGPT-4原理】
人工智能·深度学习·学习·语言模型·多模态
元宇宙时间7 小时前
Playfun即将开启大型Web3线上活动,打造沉浸式GameFi体验生态
人工智能·去中心化·区块链
开发者工具分享7 小时前
文本音频违规识别工具排行榜(12选)
人工智能·音视频
产品经理独孤虾7 小时前
人工智能大模型如何助力电商产品经理打造高效的商品工业属性画像
人工智能·机器学习·ai·大模型·产品经理·商品画像·商品工业属性
老任与码7 小时前
Spring AI Alibaba(1)——基本使用
java·人工智能·后端·springaialibaba
蹦蹦跳跳真可爱5898 小时前
Python----OpenCV(图像増强——高通滤波(索贝尔算子、沙尔算子、拉普拉斯算子),图像浮雕与特效处理)
人工智能·python·opencv·计算机视觉
雷羿 LexChien8 小时前
从 Prompt 管理到人格稳定:探索 Cursor AI 编辑器如何赋能 Prompt 工程与人格风格设计(上)
人工智能·python·llm·编辑器·prompt
两棵雪松8 小时前
如何通过向量化技术比较两段文本是否相似?
人工智能