【深度学习实战/计算机视觉实战】PyTorch框架下基于ResNet的102种花卉图像分类模型

引言

数据预处理部分：

数据增强：torchvision中transforms模块自带功能，比较实用
数据预处理：torchvision中transforms也帮我们实现好了，直接调用即可
DataLoader模块直接读取batch数据

网络模块设置：

加载预训练模型，torchvision中有很多经典网络架构，调用起来十分方便，并且可以用人家训练好的权重参数来继续训练，也就是所谓的迁移学习
需要注意的是别人训练好的任务跟咱们的可不是完全一样，需要把最后的head层改一改，一般也就是最后的全连接层，改成咱们自己的任务
训练时可以全部重头训练，也可以只训练最后咱们任务的层，因为前几层都是做特征提取的，本质任务目标是一致的

什么是预训练模型？

模型已经在大型数据集(通常是ImageNet,包含1000个类别)上训练过
学习到了通用的图像特征（边缘、纹理、形状等）
权重已经优化到可以识别各种常见物体

预训练模型的优势：

节省时间和计算资源：不需要从零开始训练
小数据集友好：即使数据量有限也能获得不错的效果
更好的泛化能力：已经学习了通用的图像特征
迁移学习的基础：可以微调以适应特定任务

微调(Fine-tuning):

在迁移学习中，我们通常保留预训练模型的大部分层
只替换和重新训练最后一层（或几层）以适应新的任务。这里我们重置最后一层全连接层
因为原始resnet的全连接层是为ImageNet的1000个类别设计的,而我们的任务可能具有不同数量的类别

网络模型保存与测试

模型保存的时候可以带有选择性，例如在验证集中如果当前效果好则保存
读取模型进行实际测试

**PS：**以下代码在 VS Code 的 Jupyter Notebook 中运行，环境为 torch2.5.1 + Python 3.9 + CUDA 11.8。请按顺序依次执行即可。文件夹结构如下：

文件夹放在百度网盘，需要自取：通过网盘分享的文件：Flowers_Classification.zip链接: https://pan.baidu.com/s/1Pyxs0ZTj1WQz36hVcSqTmQ?pwd=8e2b 提取码: 8e2b --来自百度网盘超级会员v1的分享

导入必要的库

python 复制代码

import os
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import torch
from torch import nn
import torch.optim as optim
import torchvision
from torchvision import transforms, models, datasets
import imageio
import time
import warnings
import random
import sys
import copy
import json
from PIL import Image

数据读取与预处理操作

python 复制代码

data_dir = '../flower_data/'
train_dir = data_dir + '/train'
valid_dir = data_dir + '/valid'

PS： 根据自己的文件夹结构，修改路径

制作好数据源：

data_transforms中指定了所有图像预处理操作
ImageFolder假设所有的文件按文件夹保存好，每个文件夹下面存贮同一类别的图片，文件夹的名字为分类的名字
对于训练集，我们通常使用数据增强来增加模型的泛化能力。对于验证集，我们通常只进行必要的预处理，不进行数据增强。

python 复制代码

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomRotation(45),# 随机旋转图像：角度范围为±45度
        transforms.CenterCrop(224),# 中心裁剪：从图像中心裁剪224x224的区域
        transforms.RandomHorizontalFlip(p=0.5),# 随机水平翻转：以0.5的概率水平翻转图像
        transforms.RandomVerticalFlip(p=0.5),# 随机垂直翻转：以0.5的概率垂直翻转图像
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'valid': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

batch_size = 16

# 创建数据集字典：分别为'train'和'valid'阶段创建数据集
# 使用字典推导式遍历['train', 'valid']列表
# 对每个阶段x：
#   1. os.path.join(data_dir, x): 拼接数据目录路径，如'./data/train'
#   2. datasets.ImageFolder(): 使用ImageFolder加载数据集，自动根据子文件夹名称分类
#      - 第一个参数: 数据集路径
#      - 第二个参数: 对应的数据转换（从data_transforms字典中获取）
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'valid']}

# 创建数据加载器字典：为每个数据集创建对应的DataLoader
# 使用字典推导式遍历['train', 'valid']列表
# 对每个阶段x：
#   torch.utils.data.DataLoader(): 创建数据加载器
#      - 第一个参数: 对应的数据集（从image_datasets字典中获取）
#      - batch_size: 批处理大小，一次加载多少样本
#      - shuffle: 是否在每个epoch开始时打乱数据
#        * 训练集通常设为True以增加随机性
#        * 验证集通常设为False以保持评估一致性
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True) for x in ['train', 'valid']}

# 计算数据集大小字典：获取每个数据集的样本数量
# 使用字典推导式遍历['train', 'valid']列表
# 对每个阶段x：
#   len(image_datasets[x]): 获取对应数据集的样本总数
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid']}

# 获取类别名称列表：从训练数据集中提取所有类别名称
# image_datasets['train'].classes: ImageFolder自动根据子文件夹名称生成的类别名称列表
# 例如：如果train文件夹下有'cat'、'dog'、'bird'三个子文件夹，则classes = ['cat', 'dog', 'bird']
class_names = image_datasets['train'].classes

**PS：**batch_size我设置为16，可根据实际情况写值

查看一下结构

python 复制代码

image_datasets
dataloaders
dataset_sizes

读取标签对应的花名

python 复制代码

with open('../cat_to_name.json', 'r') as f:
    cat_to_name = json.load(f)
cat_to_name

显示一个batch的图像看看

python 复制代码

# 功能：将张量形式的图像数据转换为numpy数组，并进行反标准化处理，然后使用matplotlib显示图像
def im_convert(tensor):
    """ 展示数据"""
    
    image = tensor.to("cpu").clone().detach()
    image = image.numpy().squeeze()
    image = image.transpose(1,2,0)
    image = image * np.array((0.229, 0.224, 0.225)) + np.array((0.485, 0.456, 0.406))
    image = image.clip(0, 1)

    return image

fig=plt.figure(figsize=(20, 12))
columns = 4
rows = 2

dataiter = iter(dataloaders['valid'])
inputs, classes = next(dataiter)

for idx in range (columns*rows):
    ax = fig.add_subplot(rows, columns, idx+1, xticks=[], yticks=[])
    ax.set_title(cat_to_name[str(int(class_names[classes[idx]]))])
    plt.imshow(im_convert(inputs[idx]))
plt.show()

用pytorch的nn.models中提供的模型，并且直接用训练的好权重当做初始化参数

python 复制代码

model_name = 'resnet'  #可选的比较多 ['resnet', 'alexnet', 'vgg', 'squeezenet', 'densenet', 'inception']
#是否用人家训练好的特征来做
feature_extract = True

是否用GPU训练

python 复制代码

train_on_gpu = torch.cuda.is_available()

if not train_on_gpu:
    print('CUDA is not available.  Training on CPU ...')
else:
    print('CUDA is available!  Training on GPU ...')
    
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

编写冻结所有参数的函数

python 复制代码

def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():# 冻结所有参数
            param.requires_grad = False

加载resnet模型(152层)

常见可选18、50、101、152，根据自己的实际情况选择，我选152层
第一次执行需要下载，可能会比较慢，我会提供给大家一份下载好的，可以直接放到相应路径

python 复制代码

model_ft = models.resnet152()

看下模型结构

python 复制代码

model_ft

可以看到最终输出特征是1000层，我们目标是102个类别。因为是别人训练好的模型，所以跟咱们任务是不太一样的，咱们需要把这块改成自己任务所需的结果。

resnet模型的初始化修改全连接层

python 复制代码

model_ft = models.resnet152(pretrained=True)# 加载预训练模型
set_parameter_requires_grad(model_ft, feature_extract)# 冻结所有参数
num_ftrs = model_ft.fc.in_features # 获取特征维度
# model_ft.fc 是ResNet152的最后一个全连接层,in_features 获取该层的输入特征维度

model_ft.fc = nn.Sequential(
    nn.Linear(num_ftrs, 102),	# 新线性层,将特征映射到102个类别
    nn.LogSoftmax(dim=1) 		# LogSoftmax激活
)

input_size = 224 # 设置输入尺寸

设置哪些层需要训练

python 复制代码

# GPU/CPU计算
model_ft = model_ft.to(device)

# 模型保存
filename='flowers_classification_resnet152.pth'

# 是否训练所有层
params_to_update = model_ft.parameters()
print("Params to learn:")
if feature_extract:	# 之前已赋值为True,表示直接冻住特征,不更新权重
    params_to_update = [] # 重新初始化为空列表
    for name,param in model_ft.named_parameters():# 只收集需要梯度的参数
        if param.requires_grad == True:
            params_to_update.append(param)# 将要更新的添加至列表
            print("\t",name)# 打印要更新的参数名
else:
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
            print("\t",name)

再看看模型结构

python 复制代码

model_ft

优化器、损失函数的设置

python 复制代码

# 优化器设置
optimizer_ft = optim.Adam(params_to_update, lr=1e-2)
scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)#学习率每7个epoch衰减成原来的1/10
# 最后一层已经LogSoftmax()了，所以不能nn.CrossEntropyLoss()来计算了，nn.CrossEntropyLoss()相当于logSoftmax()和nn.NLLLoss()整合
criterion = nn.NLLLoss()

定义训练模块

python 复制代码

def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, filename=filename):
    since = time.time()     # 记录开始时间
    best_acc = 0            # 最佳准确率
    model.to(device)        # 将模型移到GPU/CPU

    val_acc_history = []     # 历史验证准确率
    train_acc_history = []   # 历史训练准确率
    train_losses = []        # 训练损失
    valid_losses = []        # 验证损失
    LRs = [optimizer.param_groups[0]['lr']] # 学习率

    best_model_wts = copy.deepcopy(model.state_dict()) # 深拷贝模型参数

    # ------------------------ 训练循环 ------------------------
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 20)

        # 训练和验证
        for phase in ['train', 'valid']:
            if phase == 'train':
                model.train()  # 训练模式
            else:
                model.eval()   # 验证模式

            running_loss = 0.0
            running_corrects = 0

            # 把数据都取个遍
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 清零
                optimizer.zero_grad()
                # 只有训练的时候计算和更新梯度
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)             # 前向传播
                    loss = criterion(outputs, labels)   # 计算损失

                    _, preds = torch.max(outputs, 1)    # 获取预测类别(取最大概率索引)

                    # 训练阶段更新权重
                    if phase == 'train':
                        loss.backward()     # 反向传播，计算梯度
                        optimizer.step()    # 更新参数

                # 累计损失和正确预测数
                running_loss += loss.item() * inputs.size(0)# 0表示batch那个维度
                running_corrects += torch.sum(preds == labels.data)

            # 计算整个epoch的平均损失和准确率
            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
            
            
            time_elapsed = time.time() - since  # 统计一个epoch的耗费时间
            print('Time elapsed {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
            

            # 得到最好那次的模型
            if phase == 'valid' and epoch_acc > best_acc:
                best_acc = epoch_acc    # 更新最佳准确率
                best_model_wts = copy.deepcopy(model.state_dict())  # 保存最佳参数
                state = {
                  'state_dict': model.state_dict(),     # 模型参数
                  'best_acc': best_acc,                 # 最佳准确率
                  'optimizer' : optimizer.state_dict(), # 优化器状态(可恢复训练)
                }
                torch.save(state, '../model/{}'.format(filename))# 保存到文件
            # 验证阶段记录
            if phase == 'valid':
                val_acc_history.append(epoch_acc)
                valid_losses.append(epoch_loss)
            # 训练阶段记录
            if phase == 'train':
                train_acc_history.append(epoch_acc)
                train_losses.append(epoch_loss)

        # 记录学习率变化
        print('Optimizer learning rate : {:.7f}'.format(optimizer.param_groups[0]['lr']))
        LRs.append(optimizer.param_groups[0]['lr'])
        print()

        # 学习率衰减策略
        scheduler.step()# 学习率调度器,按照StepLR的设定,每7个epoch乘以0.1

    # ------------------------ 训练完成的处理 ------------------------
    time_elapsed = time.time() - since  # 计算总时间
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # 训练完后用最好的一次当做模型最终的结果
    model.load_state_dict(best_model_wts)
    return model, val_acc_history, train_acc_history, valid_losses, train_losses, LRs

开始训练

目前只训练了输出层

python 复制代码

model_ft, val_acc_history, train_acc_history, valid_losses, train_losses, LRs  = train_model(model_ft, dataloaders, criterion, optimizer_ft, num_epochs=20)

观察学习率衰减，确实7个epoch衰减一次

训练花了23min，观察准确率，有些过拟合，准确率还行最佳模型在最后一轮，准确率79%

现在我们全连接层已经挺优秀了，我们可以将前面的层解冻，再继续训练试试看。（可选）

这就是迁移学习中的两阶段训练策略（Two-stage Training），通常在以下场景中使用：

第一阶段：冻结特征提取层（卷积层），只训练新添加的分类层（全连接层）。
第二阶段：解冻所有层（或部分层），用较小的学习率微调整个模型。

第一阶段用较大的学习率训练新添加的分类层，快速适应新任务。第二阶段解冻所有层，用较小的学习率微调整个模型，使得预训练模型的特征提取部分也能适应新任务的数据。

**注意：**在第二阶段，由于预训练模型的参数已经很好，所以我们只进行微调，因此学习率要设置得较小，避免破坏已有的特征提取能力。

在解冻其他层后训练所有层

python 复制代码

for param in model_ft.parameters():
    param.requires_grad = True	# 解冻所有层

# 再继续训练所有的参数，这里的学习率实际会被后来加载模型给替换掉
optimizer = optim.Adam(params_to_update, lr=1e-4)
scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

# 损失函数
criterion = nn.NLLLoss()

加载之前训练好的权重参数

python 复制代码

checkpoint = torch.load('../model/{}'.format(filename)) # 加载模型
best_acc = checkpoint['best_acc']   # 加载模型参数
model_ft.load_state_dict(checkpoint['state_dict'])  # 加载模型参数
optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器状态

开始第二轮训练（训练所有层）

python 复制代码

model_ft, val_acc_history, train_acc_history, valid_losses, train_losses, LRs  = train_model(model_ft, dataloaders, criterion, optimizer, num_epochs=10)

测试网络效果

注意预处理方法要相同

现在，将编写一个函数，使用经过训练的网络进行推理。也就是说，你将传递一张图像到网络中，并预测图像中花的类别。编写一个名为"predict"的函数，获取图像和模型，然后返回前K个最可能的类以及概率。例如这样的：

python 复制代码

probs, classes = predict(image_path, model)
print(probs)
print(classes)
> [ 0.01558163  0.01541934  0.01452626  0.01443549  0.01407339]
> ['70', '3', '45', '62', '55']

加载训练好的模型

python 复制代码

# 与之前一样，先加载预训练，在进行参数替换
model_ft = models.resnet152(pretrained=True)# 加载预训练模型
set_parameter_requires_grad(model_ft, feature_extract)# 冻结所有参数
num_ftrs = model_ft.fc.in_features # 获取特征维度
# model_ft.fc 是ResNet152的最后一个全连接层,in_features 获取该层的输入特征维度

model_ft.fc = nn.Sequential(
    nn.Linear(num_ftrs, 102),	# 新线性层,将特征映射到102个类别
    nn.LogSoftmax(dim=1) 		# LogSoftmax激活
)

input_size = 224 # 设置输入尺寸

# GPU模式
model_ft = model_ft.to(device)

# 保存文件的名字
filename='flowers_classification_resnet152.pth'

# 加载模型
checkpoint = torch.load('../model/{}'.format(filename))
best_acc = checkpoint['best_acc']
model_ft.load_state_dict(checkpoint['state_dict'])# 参数替换

测试数据预处理

测试数据处理方法需要跟训练时一直才可以
crop操作的目的是保证输入的大小是一致的（224*224）
标准化操作也是必须的，用跟训练数据相同的mean和std,但是需要注意一点训练数据是在0-1上进行标准化，所以测试数据也需要先归一化
最后一点，PyTorch中颜色通道是第一个维度，跟很多工具包都不一样，需要转换

python 复制代码

def process_image(image_path):
    # 读取测试数据
    img = Image.open(image_path)
    # Resize,thumbnail方法只能进行缩小，所以进行了判断
    if img.size[0] > img.size[1]:
        img.thumbnail((10000, 256))# 宽度>高度：高度固定为256，宽度按比例调整（最多10000）
    else:
        img.thumbnail((256, 10000))# 高度>=宽度：宽度固定为256，高度按比例调整（最多10000）

    # Crop操作：从中心裁剪出224×224的区域
    # 保证所有输入图像尺寸统一，符合模型输入要求
    left_margin = (img.width-224)/2
    bottom_margin = (img.height-224)/2
    right_margin = left_margin + 224
    top_margin = bottom_margin + 224
    img = img.crop((left_margin, bottom_margin, right_margin,   
                      top_margin))
    
    # 相同的预处理方法
    img = np.array(img)/255 # 转换为numpy数组并归一化到[0,1]
    mean = np.array([0.485, 0.456, 0.406]) #ImageNet数据集均值
    std = np.array([0.229, 0.224, 0.225]) #ImageNet数据集标准差
    img = (img - mean)/std # 标准化
    
    # 注意颜色通道应该放在第一个位置
    img = img.transpose((2, 0, 1))
    
    return img

def imshow(image, ax=None, title=None):
    """展示数据"""
    if ax is None:
        fig, ax = plt.subplots()
    
    # 颜色通道还原
    image = np.array(image).transpose((1, 2, 0))
    
    # 预处理还原 反向标准化：x = (x_norm * std) + mean
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = std * image + mean
    image = np.clip(image, 0, 1)# 确保像素值在[0,1]范围内
    
    ax.imshow(image)
    ax.set_title(title)
    
    return ax

准备要预测的图片

python 复制代码

image_path = '../image_06621.jpg'
img = process_image(image_path)
imshow(img)

检查下维度信息

python 复制代码

img.shape

一旦获得了正确格式的图像，就可以编写一个函数，用模型进行预测。一种常见的做法是预测前5个左右（通常称为top- K）最可能的类别。需要计算类概率，然后找到K个最大的值。

得到一个batch的测试数据

python 复制代码

dataiter = iter(dataloaders['valid'])
images, labels = dataiter.next()

model_ft.eval()

if train_on_gpu:
    output = model_ft(images.cuda())
else:
    output = model_ft(images)

查看一下output

output表示对一个batch中每一个数据得到其属于各个类别的可能性

python 复制代码

output.shape

得到概率最大的那个

python 复制代码

_, preds_tensor = torch.max(output, 1)

preds = np.squeeze(preds_tensor.numpy()) if not train_on_gpu else np.squeeze(preds_tensor.cpu().numpy())
preds

展示预测结果

python 复制代码

fig=plt.figure(figsize=(20, 20))
columns =4
rows = 2

for idx in range (columns*rows):
    ax = fig.add_subplot(rows, columns, idx+1, xticks=[], yticks=[])
    plt.imshow(im_convert(images[idx]))
    ax.set_title("{} ({})".format(cat_to_name[str(preds[idx])], cat_to_name[str(labels[idx].item())]),
                 color=("green" if cat_to_name[str(preds[idx])]==cat_to_name[str(labels[idx].item())] else "red"))
plt.show()

**PS：**绿色为正确预测，红色为错误预测

【深度学习实战/计算机视觉实战】PyTorch框架下基于ResNet的102种花卉图像分类模型

引言

数据预处理部分：

网络模块设置：

什么是预训练模型？

预训练模型的优势：

微调(Fine-tuning):

网络模型保存与测试

导入必要的库

数据读取与预处理操作

制作好数据源：

查看一下结构

读取标签对应的花名

显示一个batch的图像看看

用pytorch的nn.models中提供的模型，并且直接用训练的好权重当做初始化参数

是否用GPU训练

编写冻结所有参数的函数

加载resnet模型(152层)

看下模型结构

resnet模型的初始化 修改全连接层

设置哪些层需要训练

再看看模型结构

优化器、损失函数的设置

定义训练模块

开始训练

在解冻其他层后训练所有层

加载之前训练好的权重参数

开始第二轮训练（训练所有层）

测试网络效果

加载训练好的模型

测试数据预处理

准备要预测的图片

检查下维度信息

得到一个batch的测试数据

查看一下output

得到概率最大的那个

展示预测结果

resnet模型的初始化修改全连接层