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

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

相关推荐
学编程就要猛5 小时前
RAG(Retrieval Augmented Generation)
人工智能·spring ai
冬奇Lab6 小时前
一天一个开源项目(第259篇):Modern Web Guidance —— Chrome 团队给 AI 编码 Agent 的「现代 Web 最佳实践」外挂
人工智能·开源
飞哥数智坊7 小时前
打造自己的 Personal Agent,从抄大肥鱼作业开始
人工智能·deepseek
对你无可奈何7 小时前
Magpie 实战:把 WorkBuddy、Qoder 的订阅模型接进一个本地池子
人工智能
Csvn7 小时前
每请求到底多少钱?把 Token 成本算成可观测指标(C01)
人工智能·aigc·agent
吴名氏.8 小时前
Java转AI Agent教程(第十三周)混合检索与重排
java·数据库·人工智能·ai·ai agent
老金带你玩AI9 小时前
惊人的一致性!Seedream 5.5实测!
人工智能
数据技术说10 小时前
在 KES AI 数据库上搭一条 NL2SQL 链路:语义、召回、权限各放哪
人工智能·sql·代码规范
jay神10 小时前
基于 YOLOv8 的手势识别系统(Web 版)
前端·人工智能·python·深度学习·yolo·计算机视觉·计算机毕业设计