学习基于pytorch的VGG图像分类 day3

注：本系列博客在于汇总CSDN的精华帖，类似自用笔记，不做学习交流，方便以后的复习回顾，博文中的引用都注明出处，并点赞收藏原博主.

目录

 VGG模型训练

        1.导入必要的库

         2.主函数部分

                 2.1使用cpu或gpu

                 2.2对数据进行预处理

                2.3 训练集部分

                 2.4索引与标签

                2.5创建数据加载器

                 2.6验证集部分

                 2.7模型的初始化

                 2.8训练部分

                2.9评估验证部分 

                 2.10主函数入口

小结

---

 VGG模型训练

        1.导入必要的库

         导入所需的库，以及导入自定义的VGG模型模版

​
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 = 32nw = 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.在模型训练时最好实时更新数据（可以更加直观的体现)