本节将介绍一个简单而典型的基于Pytorch框架的图像分类任务,所用的数据集是手写体汉字数字,模型是LeNet-5。通过该实例的学习,可以了解基于深度网络模型的计算机视觉任务的基本流程。
目录
2.1 准备
2.1.1 数据集
(1)MNIST
只要学习过深度学习相关理论的人,都一定听说过名字叫做LeNet-5模型,它是深度学习三巨头只有Yann Lecun在1998年提出的一个CNN模型(很多人认为这是第一个具有实际应用价值的CNN模型)。在当年使用该模型可以很好地完成手写体数字的识别,而该模型所处理的手写体数字数据库称为MNIST。
MNIST全称是:Mixed National Institute of Standards and Technology databas,它包含70000张手写数字的灰度图片,每一张图片包含28 *28个像素点。数据集被分为两部分,其中训练(mnist.train)集包括60000样本,测试集(mnist.test)包含10000样本。训练集又进一步封你为 55000 个样本用于训练,5000样本用于验证。下图是MNIST样本实例图。
MNIST数据集虽然经典,但也有问题。最主要的问题是,它太简单了!相对于现在动辄上百万个参数的“大”模型,MNIST数据集要小很多,且只是简单的十类问题,因此导致现有的模型在MNIST上的分类精度都超过了95%。为了更直观地观察不同算法间的性能差异,需要用一个更复杂一点的数据集,这时Fashion-MNIST出现了。
(2)Fashion-MNIST
FashionMNIST是一个替代MNIST的图像数据集。 它是由一家德国科技公司(Zalando)整理提供。FashionMNIST 的大小、格式和训练集/测试集划分与原始的 MNIST 完全一致。60000/10000 的训练测试数据划分,28*28 的灰度图片。因此,能跑MNIST数据集的代码,只需稍加改动,就可以跑新的数据集。两个数据集的不同之处主要有两点,一是虽然两者都是以灰度图像呈现的,但MNIST呈现的是数字,背景设为0,前景设为1,FashionMNIST则是真正意义的灰度数据集。二是两者内容不同,前者被分类的是手写体数字,后者则是十类衣物服饰(分别是:T恤、裤子、套头衫、连衣裙、大衣、凉鞋、衬衫、运动鞋、包、短靴),其内容的复杂程度远高于手写体数字。下图是FashionMNIST的一个图示。
网上有很多基于FashionMNIST数据集的实例,在此就不再重复介绍。
(3)CHN-MNIST
本节实例选用的是中国版MNIST,由英国纽卡斯尔大学整理并提供,我们不妨将其称为CHN-MNIST数据集。
该数据集共由100人书写,每人重复书写10遍,因此数据集样本数为1000组,每组包括15个汉字的数字,即“零、一、二、三、四、五、六、七、八、九、十、百、千、万、亿”,总样本数为15000。图像的分辨率为300*300。
2.1.2 模型
对于这样一个简单的分类任务,不需要使用太复杂的网络,前面提到的LeNet-5 足以胜任。
对于LeNet-5网络模型的介绍,网上有很多,在此不再赘述,只贴出该模型的示意图,供大家参考。
XP
2.2 代码解析
本节是重点,将结合代码,一步一步地详细介绍实现过程。
(1)载入必要的扩展库
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
由于是第一个例程,我们对所使用的扩展库详细加以介绍:
- matplotlib库:用于绘图
- numpy库:用于数值计算
- pandas库:用于数据分析
- torch库:提供Pytorch支持
- PIL库:用于图像绘制
- tqdm库:Python提供的进度条空间库
(2)设置参数
这一部分完成的是设置一些与模型训练有关的超参数。如下面代码所示:
batch_size = 32 # 批次大小
lr = 0.003 # 学习率
epochs = 10 # 迭代轮数
save_path = './best_model.pkl' # 模型保存路径
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # 设备
各个参数的功能见注释,至于各个参数数值大小对最终结果的影响,将放在后续的章节介绍。其中最后一行是自动检测是否安装了cuda,如果是,则启动gpu加速。
(3)加载数据集
定义用于数据集处理的类:
class CustomDataset(Dataset):
def __init__(self, k, l, csv_file='./data/chnmnist/chinese_mnist.csv'):
self.df = pd.read_csv(csv_file)
self.k = {'九': int(9), '十': int(10), '百': int(11), '千': int(12), '万': int(13), '亿': int(14), '零': int(0),
'一': int(1), '二': int(2), '三': int(3), '四': int(4), '五': int(5), '六': int(6), '七': int(7),
'八': int(8)}
self.target = 'character'
self.features = ['suite_id', 'sample_id', 'code', ]
self.labels = np.asarray(self.df.iloc[:, 4])
self.y = df[self.target]
self.X = df.drop(self.target, axis=1)
def __getitem__(self, idx):
single_image_label = self.labels[idx]
class_id = self.k[single_image_label]
img = Image.open(f"./data/chnmnist/input_{self.X.iloc[idx, 0]}_{self.X.iloc[idx, 1]}_{self.X.iloc[idx, 2]}.jpg")
img = np.array(img)
return img, class_id
def __len__(self):
return len(self.X)
其中,chinese_mnist.csv文件是数据文件,每个样本占一行,每行有5列,分别代表:
- 书写人序号
- 样本序号
- 对应的分类号
- 对应的数字
- 汉字数字
另一个函数完成打开图像文件的操作,样本的文件名命名方式是“input_书写人序号_样本序号_对应分类号”的方式。
还需要对数据集进行一下预处理,并分别构建训练集(对应前700个样本)和测试集(后300个样本)便于后面的训练过程
# 1.构建索引到汉字的映射字典
num2char = {int(9): '九', int(10): '十', int(11): '百',
int(12): '千', int(13): '万', int(14): '亿',
int(0): '零', int(1): '一', int(2): '二',
int(3): '三', int(4): '四', int(5): '五',
int(6): '六', int(7): '七', int(8): '八'}
# 2.读取csv处理文件
df = pd.read_csv('./chinese_mnist.csv', sep=',')
# 3.处理数据
train_df = df.groupby('value').apply(lambda x: x.sample(700, random_state=42)).reset_index(drop=True)
x_train, y_train = train_df.iloc[:, :-2], train_df.iloc[:, -2]
test_df = df.groupby('value').apply(lambda x: x.sample(300, random_state=42)).reset_index(drop=True)
x_test, y_test = test_df.iloc[:, :-2], test_df.iloc[:, -2]
加载迭代器
# 4.加载迭代器
train_dataset = CustomDataset(x_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataset = CustomDataset(x_test, y_test)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
dataiter = iter(test_loader)
images, labels = dataiter.next()
#plt.imshow(images[0], cmap='gray')
#num2char[int(labels[0])]
定义网络模型
class CNN(nn.Module):
def __init__(self, num_class=10):
super(CNN, self).__init__()
# 卷积层
self.conv1 = nn.Conv2d(1, 8, 5)
self.pool1 = nn.AvgPool2d((2, 2))
self.conv2 = nn.Conv2d(8, 16, 5)
self.pool2 = nn.AvgPool2d((2, 2))
self.conv3 = nn.Conv2d(16, 32, 5)
self.relu = nn.ReLU()
# 输出层
self.fc1 = nn.Linear(2592, 1024)
self.fc2 = nn.Linear(1024, num_class)
def forward(self, x):
# x: torch.Size([?, 1, 64, 64])
x = self.conv1(x) # torch.Size([32, 8, 60, 60])
x = self.relu(x)
x = self.pool1(x) # torch.Size([32, 8, 30, 30])
x = self.conv2(x) # torch.Size([32, 16, 26, 26])
x = self.relu(x)
x = self.pool2(x) # torch.Size([32, 16, 13, 13])
x = self.conv3(x) # torch.Size([32, 32, 9, 9])
x = self.relu(x)
x = x.flatten(start_dim=1) # torch.Size([32, 2592])
x = self.fc1(x) # torch.Size([32, 1024])
x = self.relu(x)
x = self.fc2(x) # torch.Size([32, 15])
return x
训练模型
# 5.模型训练
model = CNN(num_class=15)
model = model.to('cpu')
#model = model.to('gpu')
criterion = nn.CrossEntropyLoss() # 损失函数
optimizer = torch.optim.Adam(model.parameters(), lr=lr) # 优化器
best_acc = 0 # 最优精确率
best_model = None # 最优模型参数
for epoch in range(epochs):
model.train()
running_loss = 0 # 损失
epoch_acc = 0 # 每个epoch的准确率
epoch_acc_count = 0 # 每个epoch训练的样本数
train_count = 0 # 用于计算总的样本数,方便求准确率
train_bar = tqdm(train_loader)
for data in train_bar:
images, labels = data
optimizer.zero_grad()
output = model(images.unsqueeze(dim=1).float().to(device))
loss = criterion(output, labels.to(device))
loss.backward()
optimizer.step()
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
epochs,
loss)
# 计算每个epoch正确的个数
epoch_acc_count += (output.argmax(axis=1) == labels.view(-1)).sum()
train_count += len(images)
# 每个epoch对应的准确率
epoch_acc = epoch_acc_count / train_count
# 打印信息
print("【EPOCH: 】%s" % str(epoch + 1))
print("训练损失为%s" % str(running_loss))
print("训练精度为%s" % (str(epoch_acc.item() * 100)[:5]) + '%')
if epoch_acc > best_acc:
best_acc = epoch_acc
best_model = model.state_dict()
# 在训练结束保存最优的模型参数
if epoch == epochs - 1:
# 保存模型
torch.save(best_model, save_path)
print('Finished Training')
测试
# 6.测试图像
dataiter = iter(test_loader)
images, labels = dataiter.next()
# 8.获取一张测试图像
image = images[0].unsqueeze(dim=0).unsqueeze(dim=0).float()
model.eval()
# 9.获得输出结果
pred = num2char[model(image).argmax(axis=1).item()]
plt.imshow(images[0], cmap='gray')
print('【预测结果为:】 %s ' % pred)
下图所示是训练过程结果示意图,包括每一个epoch的损失和精度信息。经过10轮训练,最终的精度达到了98.64%。
需要说明的一点是,上述代码实际上是在cpu下运行的,调整到cuda下总是报错,我将在后面加以修订。
(未完,待续)