神经网络入门实战：（二十三）使用本地数据集进行训练和验证

使用本地数据集进行训练和验证

训练过程和用官网的数据集没什么差别，只是需要从本地加载数据集。

本地的数据集需要弄好格式，分别创建 train 和 val 两个文件夹，用来存放训练和验证数据集，两个文件夹里面是多个以类别名命名的子文件夹，其中存放相应类别的图片。样式大概是这样：

data_directory/class1/img1.jpgimg2.jpg...class2/img1.jpgimg2.jpg...class3/img1.jpgimg2.jpg...

使用 datasets.ImageFolder 和 DataLoader 加载本地数据集的代码如下：

import torch
import torchvision
from torchvision import transforms,datasets
from torch.utils.data import DataLoader# 加载本地训练数据集
train_data_dir = 'E:\\4_Data_sets\\respiratory waveform\\train'  # 数据集根目录路径
train_image_datasets = datasets.ImageFolder(train_data_dir, transform=dataclass_transform)
train_dataloader = DataLoader(train_image_datasets, batch_size=7, shuffle=True)
train_classes = train_image_datasets.classes  # 获取类别名称
# 加载本地验证数据集
test_data_dir = 'E:\\4_Data_sets\\respiratory waveform\\train'  # 数据集根目录路径
test_image_datasets = datasets.ImageFolder(test_data_dir, transform=dataclass_transform)
test_dataloader = DataLoader(test_image_datasets, batch_size=7, shuffle=True)
test_classes = test_image_datasets.classes  # 获取类别名称

完整训练代码如下：

这里以 resnet18 为例，不过模型保存在了 NN_models 文件里，并未给出。大家主要看看怎么使用本地数据集。

import torch
from torch import nn
import torchvision
from torchvision import transforms,datasets
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from NN_models import *
from PIL import Image
import time# 检查CUDA是否可用，并设置设备为 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")dataclass_transform = transforms.Compose([transforms.ToTensor(),transforms.Resize(224),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
])# 加载本地训练数据集
train_data_dir = 'E:\\4_Data_sets\\respiratory waveform\\train'  # 数据集根目录路径，自定义
train_image_datasets = datasets.ImageFolder(train_data_dir, transform=dataclass_transform)
train_dataloader = DataLoader(train_image_datasets, batch_size=7, shuffle=True)
train_classes = train_image_datasets.classes  # 获取类别名称
# 加载本地验证数据集
test_data_dir = 'E:\\4_Data_sets\\respiratory waveform\\train'  # 数据集根目录路径，自定义
test_image_datasets = datasets.ImageFolder(test_data_dir, transform=dataclass_transform)
test_dataloader = DataLoader(test_image_datasets, batch_size=7, shuffle=True)
test_classes = test_image_datasets.classes  # 获取类别名称# 创建网络模型
resnet18_Instance = resnet18(3, 3).to(device)
# 定义损失函数
loss = nn.CrossEntropyLoss()
# 定义优化器
learning_rate = 0.01
optimizer = torch.optim.SGD(resnet18_Instance.parameters(), lr=learning_rate, momentum=0.9)# 开始训练
total_train_step = 0
first_train_step = 0
total_test_step = 0
epoch_sum = 15  # 迭代次数# 添加tensorboard
writer = SummaryWriter('logs')start_time = time.time()
last_epoch_time = time.time()  # 记录开始训练的时间for i in range(epoch_sum):print("----------------第 {} 轮训练开始了----------------：".format(i + 1))# 训练步骤开始for data in train_dataloader:imgs, labels = dataimgs, labels = imgs.to(device), labels.to(device)  # 将数据和目标移动到GPUoutputs = resnet18_Instance(imgs)loss_real = loss(outputs, labels)  # 这里的损失变量 loss_real，千万别和损失函数 loss 相同，否则会报错！optimizer.zero_grad()loss_real.backward()optimizer.step()  # 更新模型参数total_train_step += 1# 表示第一轮训练结束，取每一轮的第一个batch_size来看看训练效果if total_train_step % 33 == 0:first_train_step += 1print("训练次数为：{}, loss为：{}".format(total_train_step, loss_real))  # 此训练次数非训练轮次# 每轮测试结束之后，计算训练的准确率，就拿第一个batch_size来看看准确率outputs_B = torch.argmax(outputs, dim=1)outputs_C = (labels == outputs_B).sum()  # True默认为1，False默认为0accuracy = (outputs_C.item() / len(outputs_B)) * 100accuracy = round(accuracy, 2)  # 2表示保留两位小数（四舍五入）print(f"训练正确率为：{accuracy}%")writer.add_scalar('first_batch_size', loss_real.item(), first_train_step)writer.add_scalar('total_batch_size', loss_real.item(), total_train_step)# 每训练一轮，就使用测试集看看训练效果total_test_loss = 0with torch.no_grad():for data in test_dataloader:imgs, labels = dataimgs, labels = imgs.to(device), labels.to(device)outputs = resnet18_Instance(imgs)loss_fake = loss(outputs, labels)total_test_loss += loss_fake.item()print("整体测试集上的LOSS为：{}".format(total_test_loss))one_epoch_time = time.time()  # 记录训练一次的时间one_cost_time = one_epoch_time - last_epoch_timeprint(f"训练此轮需要的时间为：{one_cost_time}")last_epoch_time = one_epoch_timeone_epoch_time = 0end_time = time.time()
total_time = end_time - start_time
print(f"训练总计需要的时间为：{total_time}")writer.close()

训练和验证结果如下：

----------------第 1 轮训练开始了----------------：
训练次数为：33, loss为：1.8306987285614014
训练正确率为：57.14%
整体测试集上的LOSS为：50.43653497334162
训练此轮需要的时间为：3.3512418270111084
----------------第 2 轮训练开始了----------------：
训练次数为：66, loss为：1.7029897492193413e-08
训练正确率为：100.0%
整体测试集上的LOSS为：36.150528442618
训练此轮需要的时间为：2.1926515102386475
----------------第 3 轮训练开始了----------------：
......
----------------第 15 轮训练开始了----------------：
训练次数为：495, loss为：1.5667459365431569e-06
训练正确率为：100.0%
整体测试集上的LOSS为：1.8008886913487672
训练此轮需要的时间为：2.1712865829467773
训练总计需要的时间为：32.333725929260254