冻结层(Freezing Layers)_PyTorch教程

冻结层(Freezing Layers)

在PyTorch中进行迁移学习时，经常会遇到需要冻结预训练模型的某些层的情况。冻结层意味着在训练过程中保持这些层的权重不变，只训练模型的新添加的层。这在训练速度和模型性能上都有很大的优势。

下面是一个详细的PyTorch教程，演示如何冻结预训练模型的某些层：

导入必要的库：

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

加载预训练模型并冻结部分层：

model = models.resnet18(pretrained=True)

# 冻结所有预训练模型的参数
for param in model.parameters():
    param.requires_grad = False

# 解冻最后的全连接层参数
model.fc.requires_grad = True

修改模型的全连接层：

# 修改全连接层的输出维度
model.fc = nn.Linear(model.fc.in_features, num_classes)

定义损失函数和优化器：

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

训练模型：

# 假设有训练数据和标签data和labels
for epoch in range(num_epochs):
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

通过以上步骤，你可以成功冻结预训练模型的部分层并训练新添加的层。这样可以加快训练速度，同时保留预训练模型的特征提取能力，从而使得模型在新任务上表现更好。希望这个教程对你有所帮助！