Pytorch进阶教学——训练一个图像分类模型（GPU）
目录
1、前言
2、数据集介绍
3、获取数据
4、创建网络
5、训练模型
6、测试模型
6.1、测试整个模型准确率
6.2、测试单张图片
1、前言
- 编写一个可以分类蚂蚁和蜜蜂图片的模型，使用数据集对卷积神经网络进行训练。训练后的模型可以对蚂蚁或蜜蜂的图片进行检测。
- 使用anaconda新建一个虚拟环境，安装好pytorch。后续缺什么包就安装什么包即可。
- 使用pycharm新建一个项目，配置好环境。
2、数据集介绍
- 使用的数据集为蚂蚁和蜜蜂的图片，分为训练集和测试集。
- 【注】数据集下载地址。
3、获取数据
- 代码中获取数据集使用的是txt文件，所以首先需要提取全部图片的地址和标签放入txt文件中。
- 下述代码为python提取全部图片地址和标签导出为txt文件的脚本。（自行修改）
-
import os # 导入os模块，用于操作文件路径等操作系统相关功能。 def get_file_name(file_path, output_file, type): # 绝对路径 path_list = os.listdir(file_path) # 列出指定路径下的所有文件和文件夹，并将结果存储在path_list中 with open(output_file, 'a') as file: for filename in path_list: all_file_path = os.path.join(file_path, filename) # 拼接路径 file.write(all_file_path + ' ' + type + '\n') if __name__ == '__main__': ants_file_path = r"D:\BaiduNetdiskWorkspace\PyTorch\image_recognition\hymenoptera_data\train\ants" bees_file_path = r"D:\BaiduNetdiskWorkspace\PyTorch\image_recognition\hymenoptera_data\train\bees" output_file = r"D:\BaiduNetdiskWorkspace\PyTorch\image_recognition\hymenoptera_data\train.txt" get_file_name(ants_file_path, output_file, 'ants') get_file_name(bees_file_path, output_file, 'bees') -
-
- 将全部地址修改为相对地址。
- 使用替换操作实现。例如：
- 使用替换操作实现。例如：
- 最后txt文件的内容如下：
- 新建一个dataset.py文件。
-
# 读取数据 import torch import torchvision.transforms as transforms from PIL import Image # 读取数据类 class MyDataset(torch.utils.data.Dataset): # 继承构建自定义数据集的基类 def __init__(self, datatxt, datatransform): datas = open(datatxt, 'r').readlines() # 按行读取，每行包含图像路径和标签 self.images = [] self.labels = [] self.transform = datatransform for data in datas: item = data.strip().split(' ') # 去除首尾空格并按空格分割 # 分别将图像路径和标签添加到self.images和self.labels列表中 self.images.append(item[0]) # 路径 self.labels.append(item[1]) # 标签 return def __len__(self): return len(self.images) # 获取数据集中的一个样本。接收一个索引item，根据索引获取对应的图像路径和标签 def __getitem__(self, item): imagepath, label = self.images[item], self.labels[item] image = Image.open(imagepath) # 打开图片 return self.transform(image), label # 返回转换后的图像和对应的标签 # 用于测试 if __name__ == '__main__': # 利用txt文件读取图片信息，txt文件包括图片路径和标签 traintxt = './hymenoptera_data/train.txt' valtxt = './hymenoptera_data/val.txt' # 图片转换形式 traindata_transfomer = transforms.Compose([ transforms.ToTensor(), # 转为Tensor格式 transforms.Resize(60), # 调整图像大小，调整为高度或宽度为60像素，另一边按比例调整 transforms.RandomCrop(48), # 裁剪图片，随机裁剪成高度和宽度均为48像素的部分 transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.RandomRotation(10), # 随机旋转 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 对图像进行归一化处理。对每个通道执行了均值为0.5、标准差为0.5的归一化操作 ]) valdata_transfomer = transforms.Compose([ transforms.ToTensor(), # 转为Tensor格 transforms.Resize(48), # 调整图像大小，调整为高度或宽度为48像素，另一边按比例调整 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) # 加载数据 traindataset = MyDataset(traintxt, traindata_transfomer) valdataset = MyDataset(valtxt, valdata_transfomer) print("测试集：" + str(traindataset.__len__())) print("训练集：" + str(valdataset.__len__()))
-
- 单独运行结果：（只用于测试）
4、创建网络
- 新建一个net.py文件。
- 其中创建了一个简单的三层卷积神经网络。
-
# 三层卷积神经网络 import torch # 卷积神经网络类 class SimpleConv3(torch.nn.Module): # 继承创建神经网络的基类 def __init__(self, classes): super(SimpleConv3, self).__init__() # 卷积层 self.conv1 = torch.nn.Conv2d(3, 16, 3, 2, 1) # 输入通道3，输出通道16，3*3的卷积核，步长2，边缘填充1 self.conv2 = torch.nn.Conv2d(16, 32, 3, 2, 1) # 输入通道16，输出通道32，3*3的卷积核，步长2，边缘填充1 self.conv3 = torch.nn.Conv2d(32, 64, 3, 2, 1) # 输入通道32，输出通道64，3*3的卷积核，步长2，边缘填充1 # 全连接层 self.fc1 = torch.nn.Linear(2304, 100) self.fc2 = torch.nn.Linear(100, classes) def forward(self, x): # 第一次卷积 x = torch.nn.functional.relu(self.conv1(x)) # relu为激活函数 # 第二次卷积 x = torch.nn.functional.relu(self.conv2(x)) # 第三次卷积 x = torch.nn.functional.relu(self.conv3(x)) # 展开成一维向量 x = x.view(x.size(0), -1) x = torch.nn.functional.relu(self.fc1(x)) x = self.fc2(x) return x # 用于测试 if __name__ == '__main__': inputs = torch.rand((1, 3, 48, 48)) # 生成一个随机的3通道、48x48大小的张量作为输入 net = SimpleConv3(2) # 二分类 output = net(inputs) print(output)
- 单独运行结果：（只用于测试）
5、训练模型
- 新建一个train.py文件。
- 其中可自行设置的参数都有标出。
-
# 训练模型 import matplotlib matplotlib.use('TkAgg') import matplotlib.pyplot as plt from dataset import MyDataset from net import SimpleConv3 import torch import torchvision.transforms as transforms from torch.optim import SGD # 优化相关 from torch.optim.lr_scheduler import StepLR # 优化相关 from sklearn import preprocessing # 处理label # 图片转换形式 traindata_transfomer = transforms.Compose([ transforms.ToTensor(), # 转为Tensor格式 transforms.Resize(60, antialias=True), # 调整图像大小，调整为高度或宽度为60像素，另一边按比例调整，antialias=True启用了抗锯齿功能 transforms.RandomCrop(48), # 裁剪图片，随机裁剪成高度和宽度均为48像素的部分 transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.RandomRotation(10), # 随机旋转 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 对图像进行归一化处理。对每个通道执行了均值为0.5、标准差为0.5的归一化操作 ]) if __name__ == '__main__': traintxt = './hymenoptera_data/train.txt' valtxt = './hymenoptera_data/val.txt' # 加载数据 traindataset = MyDataset(traintxt, traindata_transfomer) # 创建卷积神经网络 net = SimpleConv3(2) # 二分类 # 使用GPU device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net.to(device) # 测试GPU是否能使用 # print("The device is gpu later?:", next(net.parameters()).is_cuda) # print("The device is gpu,", next(net.parameters()).device) # 将数据提供给模型使用 traindataloader = torch.utils.data.DataLoader(traindataset, batch_size=128, shuffle=True, num_workers=1) # batch_size可以自行调节 # 优化器 optim = SGD(net.parameters(), lr=0.1, momentum=0.9) # 使用随机梯度下降（SGD）作为优化器，学习率0.1，动量0.9，加速梯度下降过程，lr可自行调节 criterion = torch.nn.CrossEntropyLoss() # 使用交叉熵损失作为损失函数 lr_step = StepLR(optim, step_size=200, gamma=0.1) # 学习率调度器，动态调整学习率，每200个epoch调整一次，每次调整缩小为原来的0.1倍，step_size可自行调节 epochs = 5 # 训练次数 accs = [] losss = [] # 训练循环 for epoch in range(0, epochs): batch = 0 running_acc = 0.0 # 精度 running_loss = 0.0 # 损失 for data in traindataloader: batch += 1 imputs, labels = data # 将标签从元组转换为tensor类型 labels = preprocessing.LabelEncoder().fit_transform(labels) labels = torch.as_tensor(labels) # 利用GPU训练模型 imputs = imputs.to(device) labels = labels.to(device) # 将数据输入至网络 output = net(imputs) # 计算损失 loss = criterion(output, labels) # 平均准确率 acc = float(torch.sum(labels == torch.argmax(output, 1))) / len(imputs) # 累加损失和准确率，后面会除以batch running_acc += acc running_loss += loss.data.item() optim.zero_grad() # 清空梯度 loss.backward() # 反向传播 optim.step() # 更新参数 lr_step.step() # 更新优化器的学习率 # 一次训练的精度和损失 running_acc = running_acc / batch running_loss = running_loss / batch accs.append(running_acc) losss.append(running_loss) print('epoch=' + str(epoch) + ' loss=' + str(running_loss) + ' acc=' + str(running_acc)) # 保存模型 torch.save(net, 'model.pth') # 保存模型的权重和结构 x = torch.randn(1, 3, 48, 48).to(device) # # 生成一个随机的3通道、48x48大小的张量作为输入，新建的张量也要送到GPU中 net = torch.load('model.pth') # 从保存的.pth文件中加载模型 net.train(False) # 设置模型为推理模式，意味着不会进行梯度计算或反向传播 torch.onnx.export(net, x, 'model.onnx') # 使用ONNX格式导出模型 # 接受模型net、示例输入x和导出的文件名model.onnx作为参数 # 可视化结果 fig = plt.figure() plot1, = plt.plot(range(len(accs)), accs) # 创建一个图形对象plot1，绘制accs列表中的数据 plot2, = plt.plot(range(len(losss)), losss) # 创建另一个图形对象plot2，绘制losss列表中的数据 plt.ylabel('epoch') # 设置y轴的标签为epoch plt.legend(handles=[plot1, plot2], labels=['acc', 'loss']) # 创建图例，指定图表中不同曲线的标签 plt.show() # 展示所绘制的图表
- 【注】本项目使用的是GPU训练模型。如果GPU可以获得，但是无法使用，可能是pytorch的版本不对，需要重新安装。
- 运行结果：
- 保存后的模型如下：
6、测试模型
6.1、测试整个模型准确率
- 利用测试集，测试整个模型的准确率。
- 新建一个test.py文件。
-
# 测试整个模型的准确率 import torch import torchvision.transforms as transforms from dataset import MyDataset # 您的数据集类 from sklearn import preprocessing # 处理label # 定义测试集的数据转换形式 valdata_transfomer = transforms.Compose([ transforms.ToTensor(), # 转为Tensor格式 transforms.Resize(60, antialias=True), # 调整图像大小，调整为高度或宽度为60像素，另一边按比例调整，antialias=True启用了抗锯齿功能 transforms.CenterCrop(48), # 中心裁剪图片，裁剪成高度和宽度均为48像素的部分 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 对图像进行归一化处理。对每个通道执行了均值为0.5、标准差为0.5的归一化操作 ]) if __name__ == '__main__': valtxt = './hymenoptera_data/val.txt' # 测试集数据路径 # 加载测试集数据 valdataset = MyDataset(valtxt, valdata_transfomer) # 加载已训练好的模型，利用GPU进行测试 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net = torch.load('model.pth').to(device) net.eval() # 将模型设置为评估模式，意味着不会进行梯度计算或反向传播 # 使用 DataLoader 加载测试集数据 valdataloader = torch.utils.data.DataLoader(valdataset, batch_size=1, shuffle=False) correct = 0 # 被正确预测的样本数 total = 0 # 测试样本数 # 测试模型 with torch.no_grad(): for data in valdataloader: images, labels = data # 将标签从元组转换为tensor类型 labels = preprocessing.LabelEncoder().fit_transform(labels) labels = torch.as_tensor(labels) # 利用GPU训练模型 images, labels = images.to(device), labels.to(device) outputs = net(images) # 输入图像并获取模型预测结果 _, predicted = torch.max(outputs.data, 1) # 获取预测值中最大概率的索引 total += labels.size(0) # 累计测试样本数量 correct += (predicted == labels).sum().item() # 计算正确预测的样本数量 # 计算并输出模型在测试集上的准确率 accuracy = 100 * correct / total print('Test Accuracy: {:.2f}%'.format(accuracy))
-
- 运行结果：
- 因为训练模型时只迭代了200次，所以准确率并不高。可以尝试提高训练次数，提高准确率。
6.2、测试单张图片
- 使用训练后的模型，对单张图片进行预测。
- 新建一个testone.py文件。
-
import torch from PIL import Image import torchvision.transforms as transforms # 定义图片预处理转换 image_transforms = transforms.Compose([ transforms.Resize(60, antialias=True), # 调整图像大小 transforms.CenterCrop(48), # 中心裁剪 transforms.ToTensor(), # 转为Tensor格式 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 归一化处理 ]) # 定义类别映射字典 class_mapping = { 0: "ant", 1: "bee" } # 加载已训练好的模型，利用GPU测试 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net = torch.load('model.pth').to(device) net.eval() # 将模型设置为评估模式，意味着不会进行梯度计算或反向传播 # 加载要测试的图片 image_path = './hymenoptera_data/val/bees/26589803_5ba7000313.jpg' # 图片路径 input_image = Image.open(image_path) # 加载图片 input_tensor = image_transforms(input_image).unsqueeze(0) # 对图片进行预处理转换，并增加 batch 维度 # 将输入数据移动到GPU上 input_tensor = input_tensor.to(device) # 使用模型进行预测 with torch.no_grad(): output = net(input_tensor) _, predicted = torch.max(output, 1) # 在张量中沿指定维度找到最大值及其对应的索引 # 输出预测结果 predicted_class = predicted.item() # 得到预测的标签 predicted_label = class_mapping[predicted_class] # 将标签转换为文字 print(f"The predicted class for the image is: {predicted_label}")
-
- 运行结果：
