学习基于pytorch的VGG图像分类 day3_pytorch vgg图像分类-CSDN博客


import os  
import sys  
import json  
  
import torch  
import torch.nn as nn  
from torchvision import transforms, datasets  
import torch.optim as optim  
from tqdm import tqdm  
  
from model import vgg  # 导入自定义的VGG模型模块

2.主函数部分

def main():

2.1使用cpu或gpu

# 检查是否有可用的CUDA设备，如果有则使用GPU，否则使用CPU  
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  
    print("using {} device.".format(device))

2.2对数据进行预处理

# 定义数据预处理操作，包括随机裁剪、随机水平翻转、转为Tensor格式、标准化  
    data_transform = {  
        "train": transforms.Compose([  
            transforms.RandomResizedCrop(224),  # 随机裁剪为224x224大小  
            transforms.RandomHorizontalFlip(),  # 随机水平翻转  
            transforms.ToTensor(),  # 将PIL Image或ndarray转换为torch.FloatTensor，并归一化到[0.0, 1.0]  
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化处理，减均值除标准差  
        ]),  
        "val": transforms.Compose([  
            transforms.Resize((224, 224)),  # 调整图片大小到224x224  
            transforms.ToTensor(),  
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])  # 标准化处理  
    }

2.3 训练集部分

读取脚本的路径，构建一个图像数据的跟路径（用条件判断断言路径是否存在，不存在进行报错）加载训练数据集

# 获取当前脚本的绝对路径  
    current_file_dir = os.path.dirname(os.path.abspath(__file__))  
    # 构建图像数据的根路径  
    image_root = os.path.join(current_file_dir, "image_path")  
    train_dir = os.path.join(image_root, "train")  # 训练集目录  
    # 确保训练集目录存在  
    assert os.path.exists(train_dir), "{} path does not exist.".format(train_dir)  
  
    # 使用torchvision.datasets的ImageFolder类加载训练数据集，它假设每个子文件夹的名称是其对应的类别  
    train_dataset = datasets.ImageFolder(root=train_dir, transform=data_transform["train"])  
  
    # 定义保存类别标签与索引对应关系的json文件路径  
    image_path = os.path.join(image_root)  
  
    # 计算训练集样本数量  
    train_num = len(train_dataset)

2.4索引与标签

    # 获取类别标签与索引的对应关系  
    flower_list = train_dataset.class_to_idx  
    # 反转字典，将索引映射到类别标签  
    cla_dict = {val: key for key, val in flower_list.items()} 
    # 将类别索引到标签的映射关系写入json文件  
    json_str = json.dumps(cla_dict, indent=4)  # 使用json库将字典转化为格式化字符串  
    with open('class_indices.json', 'w') as json_file:

2.5创建数据加载器

定义每个数据加载器使用的工作进程数量，若自身内存不够，可以小一点！！！

    # 定义每个数据加载器使用的工作进程数量  
    batch_size = 32
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # 取CPU核心数， batch_size（如果大于1）和8中的最小值  
    print('Using {} dataloader workers every process'.format(nw))  # 打印每个进程使用的数据加载器工作进程数  
  
    # 创建训练集数据加载器  
    train_loader = torch.utils.data.DataLoader(train_dataset,  
                                                batch_size=batch_size, shuffle=True,  
                                                num_workers=nw)

2.6验证集部分

道理同训练集。

    # 验证集目录  
    val_dir = os.path.join(image_path, "val")  # 验证集文件夹路径  
  
    # 检查验证集目录是否存在  
    assert os.path.exists(val_dir), "{} path does not exist.".format(val_dir)  
  
    # 加载验证数据集  
    validate_dataset = datasets.ImageFolder(root=val_dir, transform=data_transform["val"])  
  
    # 获取验证集样本数量  
    val_num = len(validate_dataset)  
    # 创建验证集数据加载器  
    validate_loader = torch.utils.data.DataLoader(validate_dataset,  
                                                  batch_size=batch_size, shuffle=False,  
                                                  num_workers=nw)  
  
    # 打印训练集和验证集的样本数量  
    print("using {} images for training, {} images for validation.".format(train_num, val_num))

2.7模型的初始化

对模型各个参数进行初始化，确定分类个数，训练轮数，最佳准确率，学习率

    # 初始化模型  
    model_name = "vgg16"  
    net = vgg(model_name=model_name, num_classes = 4, init_weights=True)  # 创建VGG16模型，类别数为4，并初始化权重  
    net.to(device)  # 将模型转移到指定的设备上（CPU或GPU）  
  
    # 定义损失函数和优化器  
    loss_function = nn.CrossEntropyLoss()  # 交叉熵损失函数，用于分类问题  
    optimizer = optim.Adam(net.parameters(), lr=0.0001)  # 使用Adam优化器，学习率为0.0001  
  
    # 设置训练轮数  
    epochs = 60
    # 初始化最佳准确率  
    best_acc = 0.0  
    # 设置模型保存路径  
    save_path = './{}Net.pth'.format(model_name)  
    # 计算训练步骤数  
    train_steps = len(train_loader)

2.8训练部分

训练及展示进度

    # 开始训练循环  
    for epoch in range(epochs):  
        # 将模型设置为训练模式  
        net.train()  
        # 初始化运行损失  
        running_loss = 0.0  
        # 使用tqdm库创建进度条，用于显示训练进度  
        train_bar = tqdm(train_loader, file=sys.stdout)  
        # 开始每个epoch的训练步骤循环  
        for step, data in enumerate(train_bar):  
            # 从数据加载器中获取图像和标签  
            images, labels = data  
            # 梯度清零  
            optimizer.zero_grad()  
            # 前向传播，计算输出  
            outputs = net(images.to(device))  
            # 计算损失  
            loss = loss_function(outputs, labels.to(device))  
            # 反向传播，计算梯度  
            loss.backward()  
            # 更新模型参数  
            optimizer.step()  
            # 更新运行损失（这部分代码在原始代码中被省略了，通常需要用于记录或展示）  
            # 将当前损失值添加到运行损失中  
            running_loss += loss.item()  
  
            # 打印训练过程中的统计信息  
            # 格式化字符串，显示当前epoch、总epoch数和当前损失值  
            train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,  
                                                                    epochs,  
                                                                    loss)

2.9评估验证部分

    # 验证模型性能  
        net.eval()  # 将模型设置为评估模式  
        acc = 0.0  # 初始化累积的正确预测数量  
        with torch.no_grad():  # 不计算梯度，节省计算资源  
            val_bar = tqdm(validate_loader, file=sys.stdout)  # 创建验证集的进度条  
            for val_data in val_bar:  # 遍历验证集数据  
                val_images, val_labels = val_data  # 获取图像和标签  
                outputs = net(val_images.to(device))  # 前向传播，获取模型输出  
                predict_y = torch.max(outputs, dim=1)[1]  # 获取预测类别  
                # 计算预测正确的数量，并累加到acc中  
                acc += torch.eq(predict_y, val_labels.to(device)).sum().item()  
  
        # 计算验证集的准确率  
        val_accurate = acc / val_num  
        # 打印当前epoch的训练损失和验证准确率  
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %  
            (epoch + 1, running_loss / train_steps, val_accurate))  
  
        # 如果当前验证准确率高于最佳准确率，则更新最佳准确率并保存模型状态  
        if val_accurate > best_acc:  
            best_acc = val_accurate  
            torch.save(net.state_dict(), save_path)  # 保存模型权重到指定路径  
  
        print('Finished Training')  # 训练完成，打印提示信息

2.10主函数入口

# 主函数入口  
if __name__ == '__main__':  
    main()  # 调用main函数，开始训练过程

小结

1.对内存不够的情况要降低batch_size的值，否则模型无法训练

2.在构建路径后，一定要用条件判断断言路径是否存在（这是一个好习惯）

3.在模型训练时最好实时更新数据（可以更加直观的体现)