使用pytorch,冻结resnet50前几层进行迁移学习

在PyTorch中,冻结ResNet50模型的前几层可以通过以下步骤进行:

import torch
import torchvision.models as models

# 加载预训练的ResNet50模型
model = models.resnet50(pretrained=True)

# 冻结需要保持不变的层,通常是前几个卷积层
for name, param in model.named_parameters():
    if 'conv1' in name or 'bn1' in name or 'layer1' in name or 'layer2' in name:
        param.requires_grad = False

# 修改最后一层进行微调
num_classes = 10  # 假设输出类别数为10
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

# 将模型移到GPU上(如果可用)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 编译和训练模型
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练模型
for epoch in range(num_epochs):
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs = inputs.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    # 打印每个epoch的损失值
    print(f"Epoch {
      
      epoch+1} Loss: {
      
      running_loss/len(train_loader)}")

在这个例子中,我们加载了预训练的ResNet50模型,并将指定的层参数设置为不需要梯度更新。具体来说,我们冻结了conv1bn1layer1layer2这些层的参数。然后,通过修改最后一层(全连接层)来适应自己的数据集。接下来,将模型移动到GPU上(如果可用),定义损失函数和优化器,并进行模型训练。

请根据你自己的数据集和任务适当调整代码。

猜你喜欢

转载自blog.csdn.net/weixin_45277161/article/details/132038482
今日推荐