当前位置：首页 > article >正文

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

article 2025/1/5 15:25:41

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

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

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

data_directory/
    class1/
        img1.jpg
        img2.jpg
        ...
    class2/
        img1.jpg
        img2.jpg
        ...
    class3/
        img1.jpg
        img2.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 = data
		imgs, labels = imgs.to(device), labels.to(device)  # 将数据和目标移动到GPU
		outputs = 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 += 1
			print("训练次数为：{}, 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默认为0
			accuracy = (outputs_C.item() / len(outputs_B)) * 100
			accuracy = 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 = 0
	with torch.no_grad():
		for data in test_dataloader:
			imgs, labels = data
			imgs, 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_time
	print(f"训练此轮需要的时间为：{one_cost_time}")
	last_epoch_time = one_epoch_time
	one_epoch_time = 0

end_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