基于经典网络架构训练图像分类模型-学习笔记_数据增强:torchvision中transforms模块自带功能,比较实用-CSDN博客

本文链接：https://blog.csdn.net/OliverFunction/article/details/140070729

引言：

在当今人工智能的浪潮中，深度学习的应用已经渗透到了各个领域，包括图像识别。在这篇博客中，我将分享我如何利用PyTorch，一个强大的深度学习框架，训练出一个能够识别不同类型花朵的模型。从数据准备、模型构建到训练和评估，我将带领大家一步一步地了解整个过程。无论你是深度学习的新手，还是有经验的实践者，希望这篇文章能为你提供有价值的参考和启发。

本篇博客是对唐宇迪ai的学习笔记。

环境介绍：

cuda：12.1

python：3.9.2

torch：2.1.0+cu121

TorchVision version: 0.16.0+cu121

模型训练流程

数据预处理部分：

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

网络模块设置：

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

网络模型保存与测试

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

1. 导入各类库函数和包

#导入各种库和包
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
#pip install torchvision
from torchvision import transforms, models, datasets
#https://pytorch.org/docs/stable/torchvision/index.html
import imageio
import time
import warnings
import random
import sys
import copy
import json
from PIL import Image

print(torch.__version__)
print("CUDA version:", torch.version.cuda)
print("Is CUDA available:", torch.cuda.is_available())
print("TorchVision version:", torchvision.__version__)

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

# 定义了两个文件目录，用来指向训练集和测试集（验证集）的目录

2. 数据增强和数据预处理

数据增强：通过对数据进行翻转，缩放，添加噪声等操作，实现将原始数据的多样化，增加训练数据的规模，提升训练的难度
数据预处理：通过对原始数据进行清洗，归一化，标准化等操作，使数据更适合机器学习，缩小不同类别因为数值原因而带来的学习差异

data_transforms中指定了所有图像的预处理和增强操作
ImageFolder假设所有的文件按文件夹保存好，每个文件夹下面存贮同一类别的图片，文件夹的名字为分类的名字

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomRotation(45),
        # 随机旋转，-45到45度之间随机选
        
        transforms.CenterCrop(224),
        # 从图像中心裁剪出一个 224x224 的区域
        
        transforms.RandomHorizontalFlip(p=0.5),
        # 随机水平翻转，选择一个概率概率，以50%的概论来旋转
        
        transforms.RandomVerticalFlip(p=0.5),
        # 随机垂直翻转
        
        transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),
        # 参数1为亮度，参数2为对比度，参数3为饱和度，参数4为色相
        # 亮度变化范围为 ±20%，对比度变化范围为 ±10%，饱和度变化范围为 ±10%，色相变化范围为 ±10%
        
        transforms.RandomGrayscale(p=0.025),
        # 概率转换成灰度率，3通道就是R=G=B
        
        transforms.ToTensor(),
        # 将图像转换为 PyTorch 张量，并且将像素值从 [0, 255] 范围缩放到 [0, 1] 范围
        # transforms.ToTensor()进行归一化操作，直接除以255
        
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])#均值，标准差
        # transforms.Normalize(mean, std) ，并根据公式进行标准化操作 - (x-mean)/std
        # mean 是一个包含三个值的列表 [0.485, 0.456, 0.406]，分别对应 RGB 三个通道的均值
        # 这些特定的均值和标准差值通常是根据大型数据集（例如 ImageNet）计算得来的，这些均值和标准差值反映了自然图像的普遍统计特性
        # 可以确保输入数据与模型在训练时看到的数据具有相同的分布特性，这有助于模型更好地泛化和适应新数据
        
    ]),
    '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])
        
        # 验证集只是为了用来最后的模型测试和评估，因此不需要增强，只需要同一大小即可
    ]),
}

# 定义了一个名为data_transforms的字典，用与存储数据增强和数据预处理的操作
# 通常在字典中，一个key对应一个value，当value是一个列表时，可以存储多个值

# transforms.Compose接受一个列表，其中每个元素都是一个转换操作。最终，这些操作会按顺序应用到输入图像上

dataset=datasets.ImageFolder(os.path.join(data_dir,'train'), data_transforms['train'])
print(dataset.classes)  # 输出类别名称列表
print(dataset.class_to_idx)  # 输出类别名称到索引的映射
print(len(dataset))  # 输出数据集的样本数量
print(dataset[1]) #返回第 i 个样本的图像和标签

batch_size = 8
# 设置训练的最小批次为8

image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'valid']}
# ImageFolder假设所有的文件按文件夹保存，每个文件夹下存储同一个类别的图片，文件夹名为类名
# dataset=torchvision.datasets.ImageFolder(root, transform)，其中root为图片存储的根目录，也就是各类别文件夹的上一层目录,transform会对每个图象进行预处理
# 创建一个可以直接用于深度学习模型训练的ImageFolder图像数据集对象，这是一个数据集类型

# image_datasets 是一个字典，这个字典的键分别是 'train' 和 'valid'，对应的值是两个 datasets.ImageFolder 对象
# 这个对象代表了训练数据集,它从 data_dir/train 目录中加载图像数据,并应用 data_transforms['train'] 中定义的预处理操作
# 拼接路径 path = os.path.join(data_dir, x)，data_dir/x

dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True) for x in ['train', 'valid']}
# 创建一个包含训练和验证数据加载器的字典
# torch.utils.data.DataLoader：这是 PyTorch 提供的一个类，用于包装数据集，使其能够按批次加载数据，并支持多线程数据加载和数据打乱等功能

dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid']}
# 创建一个包含训练和验证数据集大小的字典

class_names = image_datasets['train'].classes
# 获取训练数据集中的所有类别名称

ImageFolder目录展示：

with open('cat_to_name.json', 'r') as f:
    cat_to_name = json.load(f)
    
# 从名为 'cat_to_name.json' 的文件中读取数据,并将其加载到 Python 变量f
# json.load(f) 函数会将 JSON 格式的数据转换为 Python 中的字典(dictionary)对象
cat_to_name

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

model_name = 'resnet'  
# 可选的比较多 ['resnet', 'alexnet', 'vgg', 'squeezenet', 'densenet', 'inception']

feature_extract = True 
# 这表示你想使用特征提取模式，只训练模型的最后一层（输出层）
# 冻结前几层的权重参数，使它们在训练过程中保持不变，而只训练和更新最后一层（输出层）的权重参数

# 迁移学习：用在大规模数据集上预训练的模型，只微调某几层
# 因为这些预训练模型已经在大规模数据集上优化了其结构和权重，可以捕捉到数据中的重要模式，即使是不同的具体任务，这些模式也往往是有用的
# 我们可以利用已经学习到的丰富特征，只需在我们的数据集上进行少量的微调即可
# 主要是因为这些预训练模型已经更新了比较合理的权重参数，可以直接套用，只需要微调某几层即可

train_on_gpu = torch.cuda.is_available()
# 这一行代码检查系统上是否有CUDA支持的GPU
# 检查用GPU训练

if not train_on_gpu:
    print('CUDA is not available.  Training on CPU ...')
else:
    print('CUDA is available!  Training on GPU ...')

print(torch.cuda.device_count())

# 判断支持cuda的gpu数量
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 判断是否有支持cuda的gpu，如果有，则将计算设备设置为 GPU("cuda:0")
device

def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
    # 控制是否要冻结模型参数
    
        for param in model.parameters():
            param.requires_grad = False
            
# 如果feature_extracting=1，则表明冻结模型的前几层，不需要更新梯度，因此不需要更新参数
            
# model.parameters()：返回一个生成器，生成模型中所有参数（权重和偏置），不包括参数的名称。
# model.named_parameters()：返回一个生成器，生成模型中所有参数的名称和参数本身（以元组形式：(name, parameter)）
# parameter：是 torch.nn.Parameter 对象，包含参数的实际值和一些属性（如 requires_grad），表示模型中对应的可学习参数

model_ft = models.resnet18()
model_ft

# ResNet18的基本含义是，网络的基本架构是ResNet，网络的深度是18层。但是这里的网络深度指的是网络的权重层，这里包含了17个卷积层和1个全连接层
# 网络的权重层值得是参数会进行更新和学习的层

4. 将模型的输出层改成自己的

def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):
    # 定义自己的神经网络函数
    # model_name 模型名称，num_classes 分类数量，feature_extract 是否进行只进行特征提取（即冻结模型参数），use_pretrained=true 是否使用预训练模型参数
    # 特征提取，将原始数据转化为需要的特征
    
    model_ft = models.resnet18(pretrained=use_pretrained)
    # 使用预训练模型及其权重参数
    
    set_parameter_requires_grad(model_ft, feature_extract)
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_ftrs, 102)
    # 定义一个全连接层，input_features=num_ftrs,output_features=102，102为分类类别个数
    # W为102*num_ftrs，b为batch_size*102  
    # 默认 .requires_grad = True
    
    input_size = 64
    return model_ft,input_size
    # 按顺序返回

5. 设置那些层需要训练

model_ft, input_size = initialize_model(model_name, 102, feature_extract, use_pretrained=True)
# 定义了一个模型

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

filename='best.pt'
# 模型保存


params_to_update = model_ft.parameters()
# 保存所有的权重参数

print("Params to learn:")
if feature_extract:
    params_to_update = []
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
        # 判断是否需要梯度更新
        
            params_to_update.append(param)
            # 需要的话，添加入空列表
            print("\t",name)
            print(param)
else:
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
            print("\t",name)

6. 优化器设置

optimizer_ft = optim.Adam(params_to_update, lr=1e-2)
# 定义一个Adam优化器
# 优化器 optimizer 包含一个参数组和学习率

scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.1)
# 学习率每10个epoch衰减成原来的1/10
# 降低学习率可以让模型更细致地调整参数,加快最终收敛，否则模型可能会在最优解附近来回震荡,难以收敛

criterion = nn.CrossEntropyLoss()
# 定义一个损失函数

7. 训练模块

def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, filename=filename):
# 定义一个train_model函数，用来训练和测试深度学习模型
# model深度学习模型，dataloaders包含了训练和验证的数据加载器字典，criterion为损失函数，optimizer为优化器，num_epochs总训练轮数，filename保存最佳的模型名
    since = time.time()
    # 计算当前时间戳，目的是为了计算总时长
    
    best_acc = 0
    # 初始化最佳准确率为0
    
    model.to(device)
    # 将模型在GPU上运行

    val_acc_history = []
    train_acc_history = []
    # 分别记录每个epoch的验证和训练准确率
    
    train_losses = []
    valid_losses = []
    # 分别记录每个epoch的训练和验证损失
    
    LRs = [optimizer.param_groups[0]['lr']]
    # 优化器的参数组存储在 param_groups 属性中，这是一个列表，其中每个元素都是一个字典，包含该参数组的超参数和参数
    # optimizer.param_groups[0] 访问的是优化器的第一个参数组，而 ['lr'] 则获取该参数组的学习率

    best_model_wts = copy.deepcopy(model.state_dict())
    # 用于保存当前模型的权重（参数）状态的深拷贝
    # 通过 model.state_dict()，你可以获取当前模型的所有参数和状态，是一个字典
    # copy.deepcopy() 是 Python 标准库 copy 模块中的一个函数，用于创建一个对象的深拷贝；对象完全相同，但是存储在内存的不同位置

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)
        # {}中的第一个占位符会被 format() 方法中的第一个参数 epoch 替换，{}中的第二个占位符会被 format() 方法中的第二个参数 num_epochs - 1 替换

        # 训练和验证
        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 和 labels 分别代表了当前 batch 的输入数据和标签数据
                # 如果我们有一个图像分类任务,inputs就是一个 batch 的图像数据,labels就是对应的类别标签
                
                inputs = inputs.to(device)
                labels = labels.to(device)
                # 放到gpu上跑

                # 清零
                optimizer.zero_grad()
                # 只有训练的时候计算和更新梯度
                
                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)
                # loss 是当前批次的损失，loss.item() 将这个张量转换为一个 Python 数值（标量），这个损失通常是该批次的平均损失
                # inputs 是当前批次的输入数据，inputs.size(0) 获取这个批次的样本数量，即 batch_size，inputs 是一个形状为 [batch_size, channels, height, width] 的张量
                
                # loss.item() * inputs.size(0)是将批次的平均损失转换为批次的总损失
                
                running_corrects += torch.sum(preds == labels.data)
                # 计算预测正确的个数

            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            # 计算整个批次的损失，len()函数访问Dataset对象时,它会返回数据集的总样本个数
            
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
            # 计算整个批次的正确率
            
            
            time_elapsed = time.time() - since
            # 计算训练总时长
            
            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))
            # {:.4f}代表保留四位小数，{:.0f}不输出小数点
            

            # 得到最好那次的模型
            if phase == 'valid' and epoch_acc > best_acc:
            # 在验证集上进行评估
            
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                # 创建一个独立的copy
                
                state = {
                  'state_dict': model.state_dict(),
                  'best_acc': best_acc,
                  'optimizer' : optimizer.state_dict(),
                }
                #  model.state_dict(): 这是一个包含模型所有可学习参数的字典
                # 其中键是参数名称,值是对应的参数张量
                # 这些参数包括卷积层的权重和偏置,全连接层的权重和偏置,以及BatchNorm层的统计数据等。这个字典可以用于保存和加载模型的状态。
                
                #optimizer.state_dict(): 这是一个包含优化器状态的字典
                #其中包括了优化器使用的超参数(如学习率、动量等),以及每个参数的累积梯度等信息。
                
                torch.save(state, filename)
                # 将模型状态、最佳准确率和优化器状态保存到文件中
                
            if phase == 'valid':
            # 在验证阶段，记录验证准确率和损失，并使用调度器（scheduler）更新学习率
            # 因为验证在训练之后，其实也是变相的表达，在训练完之后进行学习率的调整
            
                val_acc_history.append(epoch_acc)
                valid_losses.append(epoch_loss)
                scheduler.step(epoch_loss)
                #如果达到了 step_size 指定的 epoch 数,学习率就会乘以 gamma 进行衰减
                # scheduler.step(epoch_loss)将当前epoch的验证集损失值(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()

    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)
    # 这一行代码将模型的参数设置为之前保存的最佳模型参数 best_model_wts
    
    return model, val_acc_history, train_acc_history, valid_losses, train_losses, LRs

8. 训练过程

8.1 开始训练

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

8.2 再训练所有层

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()

# 加载之前训练的最好的那个模型

checkpoint = torch.load(filename)
# 加载一个深度学习模型，用于将保存的模型参数和优化器状态恢复到当前会话中

best_acc = checkpoint['best_acc']
# 从检查点数据中提取最佳准确度（best_acc）

model_ft.load_state_dict(checkpoint['state_dict'])
#使用检查点中的 state_dict 来恢复模型的参数

optimizer.load_state_dict(checkpoint['optimizer'])

# model.state_dict()获取模型的所有状态字典，包括模型的参数（如权重和偏置）以及其他状态信息
# model.parameters()获取模型的所有可学习参数（通常是权重和偏置），是一个生成器，生成模型的参数张量

# 再次训练
model_ft, val_acc_history, train_acc_history, valid_losses, train_losses, LRs  = train_model(model_ft, dataloaders, criterion, optimizer, num_epochs=10)

9. 加载好训练的模型

# 万一中断，或者想直接加载

model_ft, input_size = initialize_model(model_name, 102, feature_extract, use_pretrained=True)

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

# 保存文件的名字
filename='best.pt'

# 加载模型
checkpoint = torch.load(filename)
best_acc = checkpoint['best_acc']
model_ft.load_state_dict(checkpoint['state_dict'])

10. 测试数据预处理

# 得到一个batch的测试数据
dataiter = iter(dataloaders['valid'])
# dataloaders会打乱数据，随机选取8个样本
# 从验证数据加载器 (dataloaders['valid']) 中创建一个迭代器 (dataiter)，用于批量加载验证及数据

images, labels = next(dataiter)
# 这行代码从迭代器中获取一个批次的数据，返回图像 (images) 和标签 (labels)

model_ft.eval()
# 代码将模型 (model_ft) 设置为评估模式。评估模式会关闭 dropout 和 batch normalization 等训练时特有的操作，从而确保模型在评估时的行为与在训练时一致

if train_on_gpu:
    output = model_ft(images.cuda())
else:
    output = model_ft(images)
# 如果 train_on_gpu 这个标志被设置为 True，表示当前的运行环境有 GPU 可用,那么就将输入图像 images 转移到 GPU 上,然后将其传入模型 model_ft 进行推理,得到输出 output

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

preds = np.squeeze(preds_tensor.numpy()) if not train_on_gpu else np.squeeze(preds_tensor.cpu().numpy())
print(preds.shape)
preds
# np.squeeze() 的作用是去除数组中的单维度条目,使其变成一维数组

print(preds_tensor)
print(labels)

11. 结果展示

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, 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()

结言：

经过一系列的探索和实验，我们成功地利用PyTorch训练出了一个高效的花朵识别模型。从数据的准备和预处理，到模型的构建、训练和最终评估，每一步都为我们提供了宝贵的经验和知识。这个项目不仅展示了深度学习在图像识别中的强大能力，也进一步巩固了我们对PyTorch框架的理解和应用。希望这篇博客能为同样热衷于深度学习和计算机视觉的你提供一些启发和帮助。

本人也是在学习AI的道路上，欢迎大家提出问题和分享答案。

欢迎大家加群一起讨论，群号：972252910