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

githubcurry2023-08-05 20:15

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

python 复制代码

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模型，并将指定的层参数设置为不需要梯度更新。具体来说，我们冻结了conv1、bn1、layer1和layer2这些层的参数。然后，通过修改最后一层（全连接层）来适应自己的数据集。接下来，将模型移动到GPU上（如果可用），定义损失函数和优化器，并进行模型训练。

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