pytorch 微调冻结

原创

mob64ca12e8d855 2024-11-13 06:40:05 ©著作权

©著作权归作者所有：来自51CTO博客作者mob64ca12e8d855的原创作品，请联系作者获取转载授权，否则将追究法律责任

PyTorch微调与冻结模型参数的入门指南

在深度学习中，微调模型是实现特定任务（如图像分类、目标检测等）的重要方式。通过在已有的预训练模型基础上进行微调和参数冻结，我们可以加速训练并提高模型效果。本文将详细介绍如何使用PyTorch实现模型微调与冻结参数，并提供具体的代码示例和详解。

流程概述

以下是实现PyTorch微调和冻结的基本流程：

步骤	描述
1. 准备数据	加载并预处理数据集
2. 加载预训练模型	从PyTorch库中加载预训练模型
3. 冻结层	冻结部分或全部层的参数，不进行更新
4. 替换输出层	根据新的分类任务替换输出层
5. 设定优化器	指定要训练的参数集，并设置优化器
6. 训练模型	进行模型训练
7. 验证模型	使用验证集对模型进行评估

详细步骤和代码实现

接下来，我们将逐步实现上述流程。

1. 准备数据

首先，我们需要准备数据集。这可以使用torchvision库来加载常见的数据集，如CIFAR-10。

import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader

# 定义数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),   # 调整图像大小
    transforms.ToTensor(),            # 转换为Tensor
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 标准化
])

# 加载训练集和验证集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=32, shuffle=True)

val_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
val_loader = DataLoader(dataset=val_dataset, batch_size=32, shuffle=False)

2. 加载预训练模型

使用PyTorch的torchvision.models加载预训练的ResNet模型。

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

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

# 打印模型结构
print(model)

3. 冻结层

我们可以通过修改模型的参数来冻结不希望更新的层。

# 冻结所有层的参数
for param in model.parameters():
    param.requires_grad = False

4. 替换输出层

根据需要的分类数替换输出层。假设我们要将输出层改为10类（CIFAR-10）。

# 替换输出层
num_classes = 10
model.fc = nn.Linear(model.fc.in_features, num_classes)

5. 设定优化器

在定义优化器时，我们只需传递需要更新的参数（即未冻结的参数）。

# 设置优化器，仅更新最后一层的参数
optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.001)

6. 训练模型

定义训练过程。

def train(model, train_loader, optimizer, criterion, device):
    model.train() # 设置模型为训练模式
    total_loss = 0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device) # 迁移到设备
        optimizer.zero_grad()                   # 清零梯度
        outputs = model(images)                 # 正向传播
        loss = criterion(outputs, labels)       # 计算损失
        loss.backward()                         # 反向传播
        optimizer.step()                       # 更新参数
        total_loss += loss.item()                # 累加损失
    return total_loss / len(train_loader)      # 返回平均损失

7. 验证模型

对于模型验证同样也需要定义一个函数。

def validate(model, val_loader, criterion, device):
    model.eval（) # 设置模型为评估模式
    total_loss = 0
    correct = 0
    with torch.no_grad(): # 关闭梯度计算
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device) # 迁移到设备
            outputs = model(images)            # 正向传播
            loss = criterion(outputs, labels)  # 计算损失
            total_loss += loss.item()          # 累加损失
            _, predicted = torch.max(outputs, 1) # 获取预测结果
            correct += (predicted == labels).sum().item() # 计算正确预测数
    accuracy = correct / len(val_loader.dataset) # 计算准确率
    return total_loss / len(val_loader), accuracy

状态图

在整个过程中，模型状态的变化可以使用状态图来描述：

stateDiagram
    [*] --> 准备数据
    准备数据 --> 加载预训练模型
    加载预训练模型 --> 冻结层
    冻结层 --> 替换输出层
    替换输出层 --> 设定优化器
    设定优化器 --> 训练模型
    训练模型 --> 验证模型
    验证模型 --> [*]

序列图

另外，整个训练过程和验证过程可以用序列图表示：

sequenceDiagram
    participant 用户
    participant 数据Loader
    participant 模型
    participant 优化器
    participant 损失函数
    用户->>数据Loader: 加载数据
    数据Loader->>模型: 传递图像和标签
    用户->>模型: 正向传播
    模型->>损失函数: 计算损失
    用户->>优化器: 清零梯度
    优化器->>模型: 反向传播
    用户->>优化器: 更新参数
    用户->>验证: 验证模型