注：本系列博客在于汇总CSDN的精华帖，类似自用笔记，不做学习交流，方便以后的复习回顾，博文中的引用都注明出处，并点赞收藏原博主.
目录
[2.3 训练集部分](#2.3 训练集部分)
VGG模型训练
1.导入必要的库
导入所需的库，以及导入自定义的VGG模型模版
python
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.主函数部分
python
def main(): 2.1使用cpu或gpu
python
# 检查是否有可用的CUDA设备，如果有则使用GPU，否则使用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device)) 2.2对数据进行预处理
python
# 定义数据预处理操作，包括随机裁剪、随机水平翻转、转为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 训练集部分
读取脚本的路径，构建一个图像数据的跟路径（用条件判断断言路径是否存在，不存在进行报错）加载训练数据集
python
# 获取当前脚本的绝对路径
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索引与标签
python
# 获取类别标签与索引的对应关系
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创建数据加载器
定义每个数据加载器使用的工作进程数量，若自身内存不够，可以小一点！！！
python
# 定义每个数据加载器使用的工作进程数量
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验证集部分
道理同训练集。
python
# 验证集目录
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模型的初始化
对模型各个参数进行初始化，确定分类个数，训练轮数，最佳准确率，学习率
python
# 初始化模型
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训练部分
训练及展示进度
python
# 开始训练循环
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评估验证部分
python
# 验证模型性能
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主函数入口
python
# 主函数入口
if __name__ == '__main__':
main() # 调用main函数，开始训练过程小结
1.对内存不够的情况要降低batch_size的值，否则模型无法训练
2.在构建路径后，一定要用条件判断断言路径是否存在（这是一个好习惯）
3.在模型训练时最好实时更新数据（可以更加直观的体现)