迁移学习的案例

from torch.utils.data import Dataset,DataLoader
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
from torch import nn
from torchvision import models
import torch.nn.functional as F

# 模型加载与参数冻结
resnet_model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
for param in resnet_model.parameters():
    print(param)
    param.requires_grad = False  # 冻结预训练权重

# 替换全连接层(适应20类输出)
in_features = resnet_model.fc.in_features
resnet_model.fc = nn.Linear(in_features, 20)

# 收集需要更新的参数(仅全连接层参数)
params_to_update = []
for param in resnet_model.parameters():
    if param.requires_grad == True:
        params_to_update.append(param)


# 数据增强与预处理
data_transforms = {
    'train':
        transforms.Compose([
            transforms.Resize([300, 300]),
            transforms.RandomRotation(45),
            transforms.CenterCrop(224),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomVerticalFlip(p=0.5),
            transforms.RandomGrayscale(p=0.1),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    'valid':
        transforms.Compose([
            transforms.Resize([224, 224]),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
}


# 自定义数据集类(此处省略具体实现,需继承Dataset并重写相关方法)
class food_dataset(Dataset):
    def __init__(self, file_path, transform=None):
        self.file_path = file_path
        self.imgs = []
        self.labels = []
        self.transform = transform
        with open(self.file_path,encoding='utf-8') as f:
            samples = [x.strip().split(' ') for x in f.readlines()]
            for img_path, label in samples:
                self.imgs.append(img_path)
                self.labels.append(label)

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, idx):
        image = Image.open(self.imgs[idx])
        if self.transform:
            image = self.transform(image)

        label = self.labels[idx]
        label = torch.from_numpy(np.array(label, dtype=np.int64))
        return image, label


# 加载训练集和验证集
training_data = food_dataset(file_path =r'E:\pythonProject3\深度学习\卷积神经网络\食物分类\train1.txt', transform=data_transforms['train'])
test_data = food_dataset(file_path=r'E:\pythonProject3\深度学习\卷积神经网络\食物分类\test1.txt', transform=data_transforms['valid'])

# 数据加载器
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)


# 设备选择
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")


model = resnet_model.to(device)

def train(dataloader,model,loss_fn,optimizer):
    model.train()
    batch_size_num = 1
    for X ,y in dataloader:
        X,y = X.to(device),y.to(device)
        pred = model(X)
        loss = loss_fn(pred,y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        loss_value = loss.item()
        if batch_size_num % 100 ==0:
            print(f"loss:{loss_value:>7f} [number:{batch_size_num}]")
        batch_size_num +=1
best_acc = 0
def test(dataloader,model,loss_fn):
    global best_acc
    size = len(dataloader.dataset)
    num_batches= len(dataloader)
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for X ,y in dataloader:
            X,y = X.to(device),y.to(device)

            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_pj_loss = test_loss / num_batches
    test_acy = correct / size * 100
    print(f"Avg loss: {test_pj_loss:>7f} \n Accuray: {test_acy:>5.2f}%")
    # 检查是否是最佳准确率
    if correct > best_acc:
        best_acc = correct
        print(f"保存最佳模型,新最佳准确率: {test_acy:>5.2f}%")
    else:
        # 打印当前最佳准确率
        best_accuracy_percent = (best_acc / size) * 100
        print(f"最佳准确率: {best_accuracy_percent:>5.2f}%")
    return test_pj_loss, test_acy
# 损失函数、优化器、学习率调度器
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params_to_update, lr=0.001)  # 仅优化全连接层参数
# optimizer = torch.optim.Adam(resnet_model.parameters(), lr=0.001)  # 若要优化所有参数,可取消此注释
# scheduler = torch.optim.lr_scheduler.stepLR(optimizer,step_size=5,gamma=0.5)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode='min',  # 指标越小越好(如损失)
    factor=0.5,  # 学习率调整倍数
    patience=3   # 多少个epoch指标无改善则调整
)


i=10
for j in range(i):
    print(f"Epoch {j+1}\n----------")
    train(train_dataloader, model,loss_fn,optimizer)
    test_loss, test_acy = test(test_dataloader, model, loss_fn)
    # scheduler.step()
    scheduler.step(test_loss)

这段代码是基于 PyTorch 实现的迁移学习(特征提取法)实战案例,核心是用预训练的 ResNet18 模型解决 20 类食物分类任务。下面我们来逐部分解析代码:

一、核心思路回顾

代码用的是「预训练模型做特征提取」方法:

  1. 加载别人训好的 ResNet18(在 ImageNet 1000 类图片上预训练),冻结主干网络参数(不让它忘记 “识别通用图像特征” 的能力);
  2. 把 ResNet18 的 “分类头”(最后一层全连接层)换成适合自己任务的 “2 分类头”(这里是 20 类食物,所以输出维度是 20);
  3. 只训练新换的全连接层参数,用少量食物数据就能快速出效果。

二、代码模块逐段解析

1. 导入依赖库(基础工具准备)
from torch.utils.data import Dataset,DataLoader  # 数据加载核心库(Dataset定义数据格式,DataLoader批量加载)
import torch  # PyTorch核心库(张量计算、模型训练)
import numpy as np  # 数值计算(处理标签转换)
from PIL import Image  # 图片读取库(加载图像文件)
from torchvision import transforms  # 图像预处理( resize、翻转、归一化等)
from torch import nn  # 神经网络层(定义全连接层、损失函数)
from torchvision import models  # 预训练模型库(直接调用ResNet18)
import torch.nn.functional as F  # 常用激活函数等(这里未直接用,预留)

作用:把后续需要的 “数据处理、模型构建、训练工具” 都导入,相当于 “备好工具箱”。

2. 模型加载与参数冻结(迁移学习核心步骤)
# 1. 加载预训练ResNet18(带默认权重,即ImageNet预训练结果)
resnet_model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

# 2. 冻结主干网络所有参数(关键!不让预训练的特征提取层“失忆”)
for param in resnet_model.parameters():
    print(param)  # (可选)打印参数,查看主干网络权重(可注释掉减少输出)
    param.requires_grad = False  # 设为False:训练时不更新这些参数

# 3. 替换全连接层(适配20类食物分类任务)
in_features = resnet_model.fc.in_features  # 获取原全连接层的输入维度(ResNet18默认是512)
resnet_model.fc = nn.Linear(in_features, 20)  # 新全连接层:输入512,输出20(对应20类)

# 4. 收集需要更新的参数(只保留新全连接层的参数,减少计算量)
params_to_update = []
for param in resnet_model.parameters():
    if param.requires_grad == True:  # 只有新全连接层的requires_grad是True
        params_to_update.append(param)

关键细节

weights=models.ResNet18_Weights.DEFAULT:自动下载预训练权重(如果本地没有),避免从头训练;

冻结参数的原因:ResNet18 的主干网络(卷积层)已经学会 “边缘、纹理、形状” 等通用图像特征,这些特征对 “食物分类” 也有用,冻结后只训分类头,既快又省数据;

替换全连接层:原 ResNet18 输出 1000 类(对应 ImageNet),我们要分 20 类,所以必须换最后一层,输入维度保持和主干网络输出一致(512)。

3. 数据增强与预处理(提升模型泛化能力)
data_transforms = {
    'train':  # 训练集预处理(加数据增强,防止过拟合)
        transforms.Compose([
            transforms.Resize([300, 300]),  # 先放大到300x300(为后续裁剪留空间)
            transforms.RandomRotation(45),  # 随机旋转(-45°~45°)
            transforms.CenterCrop(224),  # 裁剪到224x224(ResNet要求的输入尺寸)
            transforms.RandomHorizontalFlip(p=0.5),  # 50%概率水平翻转
            transforms.RandomVerticalFlip(p=0.5),  # 50%概率垂直翻转
            transforms.RandomGrayscale(p=0.1),  # 10%概率转灰度图
            transforms.ToTensor(),  # 转成Tensor(PyTorch模型只能处理Tensor)
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 归一化(用ImageNet的均值/方差,预训练模型要求)
        ]),
    'valid':  # 验证集预处理(不加增强,真实评估模型效果)
        transforms.Compose([
            transforms.Resize([224, 224]),  # 直接缩放到224x224
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
}

核心逻辑

训练集加增强:通过 “旋转、翻转” 等增加数据多样性,让模型学到更通用的特征,避免 “只认训练集里的图片,换张图就错”(过拟合);

验证集不加增强:因为验证集要测 “模型在真实场景下的表现”,不能用经过修改的图片;

归一化:预训练模型是用归一化后的 ImageNet 数据训的,我们的输入必须用相同的均值 / 方差,否则模型会 “不适应”。

4. 自定义数据集类(读取自己的食物数据)
class food_dataset(Dataset):  # 继承PyTorch的Dataset类,自定义数据格式
    def __init__(self, file_path, transform=None):
        self.file_path = file_path  # txt文件路径(存图片路径和标签)
        self.imgs = []  # 存所有图片的路径
        self.labels = []  # 存所有图片的标签(20类对应的数字)
        self.transform = transform  # 预处理函数(训练/验证用不同的)
        
        # 从txt文件读取数据(txt格式:每行是“图片路径 标签”,空格分隔)
        with open(self.file_path,encoding='utf-8') as f:
            samples = [x.strip().split(' ') for x in f.readlines()]  # 按行分割,再按空格分路径和标签
            for img_path, label in samples:
                self.imgs.append(img_path)
                self.labels.append(label)

    def __len__(self):  # 必须实现:返回数据集总数量
        return len(self.imgs)

    def __getitem__(self, idx):  # 必须实现:根据索引idx返回1个样本(图片+标签)
        # 1. 读取图片
        image = Image.open(self.imgs[idx])  # 用PIL打开图片
        if self.transform:  # 应用预处理(训练集增强/验证集常规处理)
            image = self.transform(image)

        # 2. 处理标签(转成Tensor,且类型是int64,PyTorch要求)
        label = self.labels[idx]
        label = torch.from_numpy(np.array(label, dtype=np.int64))  # 字符串标签转成int64类型的Tensor
        
        return image, label  # 返回(图片Tensor,标签Tensor)

使用前提

你需要有一个train1.txttest1.txt,格式如下(每行 1 个样本):

E:\food_data\apple1.jpg 0
E:\food_data\banana2.jpg 1
...

其中 “0、1” 是食物类别的编号(0~19,共 20 类)。

5. 加载训练 / 验证集(批量喂给模型)
# 1. 用自定义的food_dataset加载数据(指定txt路径和预处理方式)
training_data = food_dataset(file_path=r'E:\...\train1.txt', transform=data_transforms['train'])
test_data = food_dataset(file_path=r'E:\...\test1.txt', transform=data_transforms['valid'])

# 2. 用DataLoader批量加载(模型训练不能单张图喂,要批量处理提升效率)
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)  # 训练集:批量64,打乱顺序(让模型学的更全面)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)  # 验证集:批量64,打乱与否不影响(但这里也设了True,不影响结果)

关键参数

batch_size=64:每次喂给模型 64 张图(根据电脑显存调整,显存小就设 16/32);

shuffle=True(训练集):每次 epoch 打乱数据顺序,避免模型 “记顺序” 而不是学特征。

6. 设备选择(适配 GPU/CPU,加速训练)
# 优先用GPU(cuda),其次用苹果芯片的MPS,最后用CPU
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")  # 打印当前用的设备(比如“Using cuda device”)

model = resnet_model.to(device)  # 把模型放到指定设备上(必须!否则模型和数据不在一个设备会报错)

为什么要指定设备

GPU(cuda)训练速度是 CPU 的 10~100 倍,比如 ResNet18 用 CPU 训可能要几小时,GPU 只要几分钟;

注意:后续的 “数据(X、y)” 也要放到同一设备(代码里X,y = X.to(device),y.to(device)就是做这个)。

7. 训练函数(模型 “学习” 的过程)
def train(dataloader, model, loss_fn, optimizer):
    model.train()  # 设为训练模式(启用 dropout、batchnorm等训练特有的层)
    batch_size_num = 1  # 记录当前训练到第几个batch
    for X, y in dataloader:  # 循环读取每个batch的(图片X,标签y)
        # 1. 把数据放到设备上(和模型同设备)
        X, y = X.to(device), y.to(device)
        
        # 2. 前向传播:模型预测结果
        pred = model(X)  # 输入X,输出20类的概率(shape:[64,20])
        
        # 3. 计算损失(预测结果和真实标签的差距)
        loss = loss_fn(pred, y)  # 用CrossEntropyLoss(多分类任务专用)
        
        # 4. 反向传播:更新参数(只更params_to_update,即新全连接层)
        optimizer.zero_grad()  # 清空上一轮的梯度(避免累积)
        loss.backward()  # 计算梯度(从损失反向传到参数)
        optimizer.step()  # 用梯度更新参数
        
        # 5. 打印损失(每100个batch打印一次,方便监控)
        loss_value = loss.item()  # 把Tensor类型的损失转成Python数值
        if batch_size_num % 100 == 0:
            print(f"loss:{loss_value:>7f} [number:{batch_size_num}]")
        batch_size_num += 1

核心流程
前向传播(算预测)→ 算损失(找差距)→ 反向传播(算梯度)→ 更新参数(缩小差距),这是深度学习训练的核心循环。

8. 测试函数(评估模型效果,保存最佳模型)
best_acc = 0  # 记录最佳准确率(初始为0)
def test(dataloader, model, loss_fn):
    global best_acc  # 用全局变量记录最佳准确率
    size = len(dataloader.dataset)  # 验证集总样本数
    num_batches = len(dataloader)  # 验证集总batch数
    model.eval()  # 设为评估模式(禁用 dropout、固定batchnorm参数)
    test_loss = 0  # 总验证损失
    correct = 0  # 正确预测的样本数
    
    with torch.no_grad():  # 禁用梯度计算(评估时不用更新参数,节省内存和时间)
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)  # 预测
            
            # 累加损失和正确数
            test_loss += loss_fn(pred, y).item()  # 累加每个batch的损失
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()  # 统计正确数:pred.argmax(1)取概率最大的类别编号
    
    # 计算平均损失和准确率
    test_avg_loss = test_loss / num_batches  # 平均每个batch的损失
    test_acc = correct / size * 100  # 准确率(百分比)
    print(f"Avg loss: {test_avg_loss:>7f} \n Accuracy: {test_acc:>5.2f}%")
    
    # 保存最佳模型(准确率更高时更新)
    if correct > best_acc:
        best_acc = correct
        print(f"保存最佳模型,新最佳准确率: {test_acc:>5.2f}%")
        # (可选)这里可以加模型保存代码:torch.save(model.state_dict(), "best_model.pth")
    else:
        # 打印当前最佳准确率
        best_acc_percent = (best_acc / size) * 100
        print(f"最佳准确率: {best_acc_percent:>5.2f}%")
    
    return test_avg_loss, test_acc

关键细节

model.eval()with torch.no_grad():评估时必须加,否则模型状态不对,准确率计算不准;

pred.argmax(1):取每个样本预测概率中最大的那个类别(比如 pred 是 [64,20],argmax (1) 后是 [64],每个元素是 0~19 的类别编号);

保存最佳模型:避免训练后期过拟合导致准确率下降,只保留表现最好的模型。

9. 配置训练工具(损失函数、优化器、学习率调度器)
# 1. 损失函数:多分类任务用CrossEntropyLoss(自带softmax,不用自己加)
loss_fn = nn.CrossEntropyLoss()

# 2. 优化器:只优化需要更新的参数(params_to_update,即新全连接层)
optimizer = torch.optim.Adam(params_to_update, lr=0.001)  # Adam是常用优化器,学习率0.001(可调整)
# (注释掉的代码)如果要做“微调(Fine-tuning)”,就优化所有参数:
# optimizer = torch.optim.Adam(resnet_model.parameters(), lr=0.001)

# 3. 学习率调度器:根据验证损失调整学习率(避免后期学习率太大导致震荡)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode='min',  # 指标“越小越好”(这里是验证损失)
    factor=0.5,  # 学习率调整倍数:如果指标没改善,就乘以0.5(比如0.001→0.0005)
    patience=3   # 连续3个epoch指标没改善,才调整学习率
)
# (注释掉的代码)固定步长调度器:每5个epoch学习率乘以0.5
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

调度器的作用

训练初期用较大的学习率(0.001)快速接近最优解;

后期如果验证损失不再下降,减小学习率(乘以 0.5),让模型更精细地调整参数,避免 “跑过头”。

10. 启动训练循环(控制训练轮次,调用训练 / 测试函数)
i = 10  # 训练10个epoch(1个epoch=把训练集完整过一遍)
for j in range(i):
    print(f"Epoch {j+1}\n----------")  # 打印当前是第几个epoch
    train(train_dataloader, model, loss_fn, optimizer)  # 训练一轮
    test_loss, test_acc = test(test_dataloader, model, loss_fn)  # 测试一轮
    scheduler.step(test_loss)  # 根据验证损失调整学习率(对应ReduceLROnPlateau)
    # scheduler.step()  # 如果用StepLR,就用这行(不用传指标)

epoch 的选择

10 个 epoch 是比较合理的初始值,可根据效果调整:如果验证准确率还在上升,就增加 epoch;如果准确率下降(过拟合),就减少 epoch 或加正则化。

三、代码整体逻辑总结(从输入到输出)

  1. 数据端:从 txt 读取食物图片路径和标签→自定义 Dataset 处理图片 + 预处理→DataLoader 批量加载;
  2. 模型端:加载预训练 ResNet18→冻结主干网络→替换 20 类全连接层;
  3. 训练端:按 epoch 循环→每个 epoch 先训练(更新分类头参数)→再测试(评估准确率)→根据验证损失调整学习率;
  4. 目标:用迁移学习快速训练出能分 20 类食物的模型,且避免过拟合、节省计算资源。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值