写点什么

【PyTorch 卷积】实战自定义的图片归类

作者:北桥苏
  • 2023-11-02
    广东
  • 本文字数:4126 字

    阅读完需:约 14 分钟

前言

​ 卷积神经网络是一类包含卷积计算且具有深度结构的前馈神经网络,是深度学习的代表算法之一,它通过卷积层、池化层、全连接层等结构,可以有效地处理如时间序列和图片数据等。关于卷积的概念网络上也比较多,这里就不一一描述了。实战为主当然要从实际问题出发,用代码的方式加深印象。在写代码前,我先说一下为什么我要写这篇文章?


​ 之前我也用 Tensorflow.js 跟着别人试过图片分类,虽然结果是有了,但是对代码的理解和印象并不深刻。后来由于工作业务原因才接触 PyTorch,发现这个框架更好上手,整一圈后就想用这个把之前用得图片也实现一下分类。开始也是看文章实现,但是网上大部分都是用 MNIST 数据集实现的手写字识别,而业务中有时就是一些指定的不规则小众图片识别,所以下面就简单实现一个自定义的图片集归类。


流程

  • 根据自己的定义,收集图片并归类

  • 读取图片数据和归类标签,保存数据集

  • 固定图片大小 (会变形),归一化转张量

  • 定义超参数,损失函数和优化器等

  • 炼丹,重复查看损失值准确率等指标

  • 保存模型参数,加载测试图片分类效果


环境

  • Python 3.8

  • Torch 1.9.0

  • Pillow 10.0

  • Torchvision

  • Numpy

  • Pandas

  • Matplotlib

编码

​ 写代码前已经把需要的图片做好了分类,上面的依赖包也已经安装完毕。由于只是演示这里没有用预训练模型(ResNet、VGG),因为训练时要用的是 Tensor,所以需要先读取文件夹内的图片先转化为 PIL 的对象数据或 Numpy 数据,然后可以对图片进行调整,最后全都转成 Tensor(也可以跳过 PIL 直接转张量)。这里需要注意的是对灰彩图片通道,不同尺寸图的统一处理,就是灰色图的单通道要通过复制的方式创建三个通道,所以图片设置一样的像素大小。因为在卷积网络中,输入的通道数和输入大小要一致,不然可能在训练中报错。


图片数据生成

​ 这里就是遍历各个分类文件夹的图片转换为对象信息数据,和提取所有分类,分别保存到指定位置,当然也可以在这里划分训练数据,校验数据,测试数据,需要的可以扩展这里就跳过了。


# -*- coding: utf-8 -*-import osimport pickle as pklimport pandas as pdfrom PIL import Image
all_cate = []data_set = []directory = "./data/train"for index, data in enumerate(os.walk(directory)): root, dirs, files = data
if index == 0: all_cate += dirs else: sorted(all_cate)
root_names = root.split("\\") dir_name = root_names[-1]
for img in files: img_path = root + "\\" + img img_np = Image.open(img_path) dict = {} dict['img_np'] = img_np dict['label'] = all_cate.index(dir_name) + 1 data_set.append(dict)
# 字典转DataFramedf = pd.DataFrame(data_set)pkl.dump(df, open('data/train_dataset.p', 'wb'))open("data/all_cate.txt", encoding="utf-8", mode="w+").write("\n".join(all_cate))
print("存档数据成功~")
复制代码

批量数据集标准化

​ 这里是读取序列化的图片信息,对所有图片统一像素 (一般配置电脑最好在 100px 以内,不然会很卡) 并标准归一化后,转换为 Tensor。然后判断图片通道数,如果是灰色图,可以复制张量三次以创建三个通道,最后通过 torch 的 DataLoader 在训练前完成数据集的加载。


# -*- coding: utf-8 -*-import torchfrom torchvision import transformsimport pickle as pklfrom torch.utils.data import Dataset
class DataSet(Dataset):
def __init__(self, pkl_file): df = pkl.load(open(pkl_file, 'rb')) self.dataFrame = df
def __len__(self): return len(self.dataFrame)
def __getitem__(self, item):
img_np = self.dataFrame.iloc[item, 0] label = self.dataFrame.iloc[item, 1]
transform = transforms.Compose([ transforms.Resize((100, 100)), # 根据需要调整图像大小 transforms.ToTensor(), transforms.Normalize([0.5], [0.5]) # 标准归一化, p1.均值 p2.方差 ]) image_tensor = transform(img_np)
if image_tensor.shape[0] == 1: image_tensor = image_tensor.repeat(3, 1, 1)
res = { 'img_tensor': image_tensor, 'label': torch.LongTensor([label-1]) # 需要实际的索引值 }
return res
复制代码

神经网络模型

​ 这里创建的是卷积神经网络,接收 3 通道,第一层卷积层卷积核 3x3,输出 25 维张量,通过批标准化(BatchNorm2d)进行归一化处理,最后通过 ReLU 激活函数进行非线性变换。第一层池化使用 2x2 的最大池化操作对卷积后的特征图进行下采样。第二层也是卷积和对应的池化,最后是全连接层。将经过池化的特征图展平,然后通过一个有 1024 个神经元的全连接层,再通过 ReLU 激活函数进行非线性变换。之后是一个有 128 个神经元的全连接层,最后再通过 ReLU 激活函数进行非线性变换,输出 5 个神经元代表分类的概率分布。


# -*- coding: utf-8 -*-import torch.nn as nnimport torchimport mathimport torch.functional as F
class CNN(nn.Module):
def __init__(self): super(CNN, self).__init__()
self.layer1 = nn.Sequential( nn.Conv2d(3, 25, kernel_size=3), nn.BatchNorm2d(25), nn.ReLU(inplace=True) )
self.layer2 = nn.Sequential( nn.MaxPool2d(kernel_size=2, stride=2) )
self.layer3 = nn.Sequential( nn.Conv2d(25, 50, kernel_size=3), nn.BatchNorm2d(50), nn.ReLU(inplace=True) )
self.layer4 = nn.Sequential( nn.MaxPool2d(kernel_size=2, stride=2) )
self.fc = nn.Sequential( nn.Linear(50 * 23 * 23, 1024), nn.ReLU(inplace=True), nn.Linear(1024, 128), nn.ReLU(inplace=True), nn.Linear(128, 5) )
def forward(self, x): x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = x.view(x.size(0), -1)
x = self.fc(x)
return x
复制代码

开始训练

# -*- coding:utf-8 -*-import torchfrom torch.utils.data import DataLoaderimport matplotlib.pyplot as pltfrom data_set import DataSetfrom torch.autograd import Variablefrom utils import *import cnnimport torch.nn as nnimport numpy as npimport torch.optim as optim
# 定义超参数batch_size = 1learning_rate = 0.02num_epoches = 1
# 加载图片tensor训练集tain_dataset = DataSet("data/train_dataset.p")train_loader = DataLoader(tain_dataset, batch_size=batch_size, shuffle=True)
model = cnn.CNN()
# 定义损失函数和优化器criterion = nn.CrossEntropyLoss()optimizer = optim.SGD(model.parameters(), lr=learning_rate)
# 训练模型train_loses = []records = []for i in range(num_epoches): for ii, data in enumerate(train_loader): img = data['img_tensor'] label = data['label'].view(-1)
optimizer.zero_grad() out = model(img) loss = criterion(out, label) train_loses.append(loss.data.item()) loss.backward() optimizer.step()
if ii % 50 == 0: print('epoch: {}, loop: {}, loss: {:.4}'.format(i, ii, np.mean(train_loses)))
records.append([np.mean(train_loses)])
# 绘制模型的损失,准确率走势图train_loss = [data[0] for data in records]plt.plot(train_loss, label = 'Train Loss')plt.xlabel('Steps')plt.ylabel('Loss')plt.legend()plt.show()
# 模型评估(略)# model.eval()
# 模型保存torch.save(model, 'params/cnn_imgs_02.pkl')
复制代码

模型检测

​ 训练完成保存参数到本地,下面就是将加载进的参数来测试其他图片的分类效果,同样的也是将指定图片和训练时一样的转换操作,最后将预测结果取出最大分布索引值,根据索引就可以匹配出分类名称了。另一个是工具函数,将 tensor 格式的图片在预测结果后显示在 pyplot 中。


# -*- coding:utf-8 -*-import torchfrom torch.utils.data import DataLoaderimport matplotlib.pyplot as pltfrom data_set import DataSetfrom utils import *import torchvisionfrom PIL import Imagefrom torchvision import transformsimport cnn

def imshow(img): img = img / 2 + 0.5 npimg = img.numpy() plt.imshow(np.transpose(npimg, (1, 2, 0))) plt.show()
model = torch.load("params/cnn_imgs_02.pkl")
img_path= "imgs/05.jpg"img_np = Image.open(img_path)transform = transforms.Compose([ transforms.Resize((100, 100)), transforms.ToTensor(), transforms.Normalize([0.5], [0.5]) ])image_tensor = transform(img_np)
# 如果是灰度图片if image_tensor.shape[0] == 1: image_tensor = image_tensor.repeat(3, 1, 1)
image_tensor = image_tensor.view(-1, 3, 100, 100)
predict = model(image_tensor)indices = torch.max(predict, 1)[1].item()
all_cate = []for line in open("data/all_cate.txt", encoding="utf-8", mode="r"): all_cate.append(line.strip())
cate_name = ""try: cate_name = all_cate[indices]except ValueError: cate_name = "未知"
print("识别结果是:", cate_name)# imshow(torchvision.utils.make_grid(image_tensor))# 原图显示img_np.show()exit()
复制代码



用户头像

北桥苏

关注

公众号:ZERO开发 2023-05-08 加入

专注后端实战技术分享,不限于PHP,Python,JavaScript, Java等语言,致力于给猿友们提供有价值,有干货的内容。

评论

发布
暂无评论
【PyTorch 卷积】实战自定义的图片归类_卷积神经网络_北桥苏_InfoQ写作社区