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

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

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

相关推荐
海兰1 天前
【游戏】迷雾镇(Mist Town)AI 沙箱游戏详细设计与部署指南(附源代码)
人工智能·游戏
小赖同学啊1 天前
智能连接器集群化高可用生产方案
linux·运维·人工智能
ZStack开发者社区1 天前
基于AI Agent的ZCF API文档全链路自动化
运维·人工智能·自动化
沈麽鬼1 天前
别瞎用AI写代码!90%开发者都搞错了AI编程的底层逻辑
人工智能·ai编程·trae
小陈爱编程1 天前
我终于把 Codex 的 API 配置理顺了:从踩坑到跑通
人工智能
不爱洗脚的小滕1 天前
【Agent】如何为 AI Agent 设计高可用的 Tools
人工智能·aigc·ai编程·rag
姗姗来迟了1 天前
前端传图片给多模态 Agent:压缩、预览、格式那些破事
人工智能
Sam09271 天前
Spec Coding 和 Vibe Coding 的区别:AI Coding 从感觉驱动到规格驱动
人工智能·ai
Kobebryant-Manba1 天前
学习RNN(简洁实现)
人工智能·rnn·学习
德迅--文琪1 天前
当前 2026 年 AI 狂潮时代,抗 DDoS 产品公司品牌推荐
人工智能·ddos