#AI挑战营第一站#手把手教你训练一个基于pytorch的手写数字识别模型
[复制链接]
手把手教你训练一个基于pytorch的手写数字识别模型
首先我们可以使用网上流传最广的MNIST数据集，这里有两种下载方式，两种下载方式效果一样，可以自由选择，如果程序内下载太慢就可以直接拷贝别人下载好的或者自己去官网下载一个数据包：
一、下载数据集：
1、官网下载
找到图中框选的四个下载链接，如果打开失败就多打开几次，可能是访问的人太多了导致服务器卡顿。
其中前两个文件时训练的图像和标签，后两个文件时测试集的图像和标签，都下载等会训练模型要用的。
2、程序内部自动下载（推荐）
程序内部下载就比较省心，只需要程序内的添加一个获取数据集的函数即可：
这个函数需要依赖import torchvision.datasets包，因此别忘了导入。
import torchvision.datasets
def get_data_loader(dat_path, bat_size, trans, to_train=False):
"""
获取MNIST数据集的数据加载器。
参数:
- dat_path: 数据集的路径。
- bat_size: 批量大小。
- trans: 数据转换器，用于数据预处理。
- to_train: 是否为训练模式，默认为False，即为测试模式。
返回值:
- dat_set: MNIST数据集。
- dat_loader: 数据加载器，用于遍历数据集。
"""
# 加载MNIST数据集，根据to_train参数决定是训练集还是测试集
dat_set = torchvision.datasets.MNIST(root=dat_path, train=to_train, transform=trans, download=True)
# 根据是否为训练模式，配置不同的数据加载器参数
if to_train is True:
dat_loader = torch.utils.data.DataLoader(dat_set, batch_size=bat_size, shuffle=True)
else:
dat_loader = torch.utils.data.DataLoader(dat_set, batch_size=bat_size)
return dat_set, dat_loader
如果你只想训练一个模型不想了解代码的细节的话可以跳过这一部分，直接复制上面代码即可。
也有人可能会问，这里没有下载链接啊，哪里看出来有下载的功能了，这就是我们起那么导入那个包的目的了，注意到我们有这么一行代码，我们设置download=true后他就会默认下载。
# 加载MNIST数据集，根据to_train参数决定是训练集还是测试集
dat_set = torchvision.datasets.MNIST(root=dat_path, train=to_train, transform=trans, download=True)
在pycharm中，我们鼠标点击MNIST，然后再按下Ctrl+B，即可跳转到它的初始化代码：
我们可以看到他有一个自带的下载函数：
再按Ctrl+B跳转到他的下载函数内，可以看到，它提示如果没有就会下载数据集：
我们看他下面下载文件的地方，点击self.mirrors然后跳转：
我们可以看到，resources里面的四个文件名和我们前面官网看到的一模一样，只是前面的镜像源不一样，所以不用担心下载的数据集不同，我个人喜欢直接程序里面下载，这样他下载好了后就会直接运行，也不需要我再去调整路径什么的，比较方便。
二、设置模型参数
下载好数据集后我们就应该设置模型参数了，学过神经网络的都知道，一般的神经网络看起来比较简单，但是如果自己写可能代码比较多，不过我们使用pytorch即可节省大量代码并且结构清晰也容易理解。
下面就是我们的普通的一个CNN，包括三个中间层一个全连接层，其中前两个中间层我们都是采用卷积加池化的配合中间加一个激活函数，这是默认搭配，经典的CNN网络都是卷积后面加池化层，卷积的功能可以理解为提取出目标中的特征，池化层的功能可以理解为将这些特征组合起来。我们的输入图像是28*28，卷积核采用3*3即可，也可以改成5*5或者7*7，想试试都可以试试，但是太大可能会导致模型效果不好，而且如果卷积核太大会导致卷积出来的结果图片像素过低无法进行下一次卷积操作。再定义一个前向传播的函数将这些层连接起来即可：
class CNN(nn.Module):
"""
CNN类：定义了一个基于卷积神经网络的模型。
"""
def __init__(self):
"""
初始化CNN模型，包括三个卷积层和一个全连接层。
"""
super(CNN, self).__init__()
# 第一个中间层：1个输入通道，16个输出通道，卷积核大小为3x3
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1),
nn.ReLU(), # 使用ReLU激活函数
nn.MaxPool2d(kernel_size=2, stride=2), # 使用2x2的最大池化层
)
# 第二个中间层：16个输入通道，32个输出通道，卷积核大小为3x3
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
# 第三个中间层：32个输入通道，64个输出通道，卷积核大小为3x3
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
)
# 全连接层：首先将卷积层输出展平，然后通过两个线性层降低维度，最后输出10个类别的概率
self.fullyConnected = nn.Sequential(
nn.Flatten(),
nn.Linear(in_features=7 * 7 * 64, out_features=128),
nn.ReLU(),
nn.Linear(in_features=128, out_features=10),
)
def forward(self, img):
"""
定义CNN的前向传播路径。
参数：
img -- 输入图像张量
返回值：
output -- 经过CNN模型处理后的输出张量
"""
output = self.conv1(img) # 应用第一个卷积层
output = self.conv2(output) # 应用第二个卷积层
output = self.conv3(output) # 应用第三个卷积层
output = self.fullyConnected(output) # 应用全连接层
return output
三、获取当前设备
了解过机器学习的都知道cuda吧，这是Nvidia公司为他们的显卡研究出来的专门用来加速类似于AI训练的东西，有他在我们的训练时间会大大降低，但是也不乏有部分人的电脑是没有Nvidia显卡甚至没有显卡（跟我一样可怜……），因此为了兼顾有显卡和没有显卡的用户，代码里面一般都会考虑到这个问题，所以我们会检测是否有cuda，有就用，没有就用cpu:
def get_device():
if torch.cuda.is_available():
train_device = torch.device('cuda')
else:
train_device = torch.device('cpu')
return train_device
四、编写训练函数
万事俱备就应该开始炼丹了家人们，定义一个训练的函数是为了方便我们后面的调用，将我们要用到的东西放进去即可，然后按照标准的网络训练过程编写代码即可，这些照抄就行了，中间输出一些信息来让我们看到训练的过程，防止程序死机了我们还发现不了。
def train(network, dat_loader, device, epos, loss_function, optimizer):
"""
对给定的网络进行训练。
参数:
- network: 待训练的网络模型。
- dat_loader: 数据加载器，用于批量加载训练数据。
- device: 指定训练过程使用的设备（如CPU或GPU）。
- epos: 训练的总轮数。
- loss_function: 指定的损失函数。
- optimizer: 优化器，用于更新网络参数。
返回值:
- 训练后的网络模型。
"""
for epoch in range(1, epos + 1):
network.train(mode=True) # 将网络设置为训练模式
for idx, (train_img, train_label) in enumerate(dat_loader):
train_img = train_img.to(device) # 将训练图像数据移动到指定设备
train_label = train_label.to(device) # 将训练标签移动到指定设备
outputs = network(train_img) # 通过网络计算输出
optimizer.zero_grad() # 清除之前的梯度
loss = loss_function(outputs, train_label) # 计算损失
loss.backward() # 反向传播计算梯度
optimizer.step() # 使用优化器更新网络参数
# 每隔100个批次打印一次训练状态
if idx % 100 == 0:
cnt = idx * len(train_img) + (epoch - 1) * len(dat_loader.dataset)
print('epoch: {}, [{}/{}({:.0f}%)], loss: {:.6f}'.format(epoch,
idx * len(train_img),
len(dat_loader.dataset),
(100 * cnt) / (
len(dat_loader.dataset) * epos),
loss.item()))
print('------------------------------------------------')
print('Training ended.')
return network
五、编写测试函数
跟训练函数类似，只不过这个是用来评估咱们刚练出来的模型效果如何的，也是套模板即可：
def test(network, dat_loader, device, loss_function):
"""
测试给定的网络模型。
参数:
- network: 训练好的神经网络模型。
- dat_loader: 测试数据集的数据加载器。
- device: 指定运行设备，如"cpu"或"cuda:0"。
- loss_function: 用于计算损失的函数。
返回值:
- 无返回值，但会打印测试集的损失、正确率等信息。
"""
test_loss_avg, correct, total = 0, 0, 0
test_loss = []
network.train(mode=False) # 将网络设置为评估模式
with torch.no_grad(): # 禁止计算梯度，以减少内存消耗
for idx, (test_img, test_label) in enumerate(dat_loader):
test_img = test_img.to(device) # 将测试图像数据移动到指定设备
test_label = test_label.to(device) # 将测试标签移动到指定设备
total += test_label.size(0) # 统计测试样本总数
outputs = network(test_img) # 通过网络获得输出结果
loss = loss_function(outputs, test_label) # 计算损失
test_loss.append(loss.item()) # 记录当前样本的损失
predictions = torch.argmax(outputs, dim=1) # 获取预测标签
correct += torch.sum(predictions == test_label) # 统计正确预测的数量
test_loss_avg = np.average(test_loss) # 计算测试损失的平均值
# 打印测试结果: 总数、正确数、准确率和平均损失
print('Total: {}, Correct: {}, Accuracy: {:.2f}%, AverageLoss: {:.6f}'.format(total, correct,
六、编写显示代码
我们第一步下载了数据集，应该都发现了数据集怎么是ubyte结尾，也没法查看啊，那是因为他们采用的 .idx1-ubyte和.idx3-ubyte 格式的文件。这是一种IDX数据格式，我们人类一般没法看懂这个文件，但是可以通过程序来显示，从网上找一个别人写好的显示这个格式的代码然后改改就能用：
def show_part_of_image(dat_loader, row, col):
"""
展示数据加载器中部分图像的网格视图。
参数:
- dat_loader: 数据加载器，应包含图像和对应的标签。
- row: 每行展示的图像数量。
- col: 每列展示的图像数量。
返回值:
- 无
"""
# 从数据加载器中获取第一个样本
iteration = enumerate(dat_loader)
idx, (exam_img, exam_label) = next(iteration)
# 创建一个图形窗口
fig = plt.figure(num=1)
# 在图形窗口中布局子图
for i in range(row * col):
plt.subplot(row, col, i + 1)
plt.tight_layout()
# 显示图像
plt.imshow(exam_img[i][0], cmap='gray', interpolation='none')
plt.title('Number: {}'.format(exam_label[i])) # 设置图像标题为对应的标签
plt.xticks([]) # 移除x轴刻度
plt.yticks([]) # 移除y轴刻度
plt.show() # 显示图像
既然都写了一个显示数据集的函数了，那我们结果当然也要显示出来看一下了：
def show_part_of_test_result(network, dat_loader, row, col):
"""
展示测试结果的一部分。
参数:
- network: 训练好的神经网络模型。
- dat_loader: 测试数据集的加载器。
- row: 展示图片的行数。
- col: 展示图片的列数。
返回值:
- 无。
"""
# 初始化测试数据的迭代器
iteration = enumerate(dat_loader)
# 获取第一批测试数据
idx, (exam_img, exam_label) = next(iteration)
# 在不计算梯度的情况下通过模型预测输出
with torch.no_grad():
outputs = network(exam_img)
# 创建一个图形用于展示图片
fig = plt.figure()
# 遍历要展示的图片并展示在图形上
for i in range(row * col):
plt.subplot(row, col, i + 1) # 创建子图
plt.tight_layout() # 调整子图间的布局以避免重叠
plt.imshow(exam_img[i][0], cmap='gray', interpolation='none') # 展示图片
plt.title('Number: {}, Prediction: {}'.format(
exam_label[i], outputs.data.max(1, keepdim=True)[1][i].item() # 标题包括真实标签和预测标签
))
plt.xticks([]) # 移除x轴标签
plt.yticks([]) # 移除y轴标签
plt.show() # 展示图形
七、所有代码展示
上面已经把代码都介绍一遍了，我相信大家应该也对模型的训练有个大致的思路了，接下来就是将所有代码整合到一起然后编写一个主函数去进行训练测试了，这里就不多说了，大家直接看代码然后就可以开始炼丹了：
import torch
import torch.nn as nn
import torchvision.datasets
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
class CNN(nn.Module):
"""
CNN类：定义了一个基于卷积神经网络的模型。
"""
def __init__(self):
"""
初始化CNN模型，包括三个卷积层和一个全连接层。
"""
super(CNN, self).__init__()
# 第一个中间层：1个输入通道，16个输出通道，卷积核大小为3x3
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1),
nn.ReLU(), # 使用ReLU激活函数
nn.MaxPool2d(kernel_size=2, stride=2), # 使用2x2的最大池化层
)
# 第二个中间层：16个输入通道，32个输出通道，卷积核大小为3x3
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
# 第三个中间层：32个输入通道，64个输出通道，卷积核大小为3x3
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
)
# 全连接层：首先将卷积层输出展平，然后通过两个线性层降低维度，最后输出10个类别的概率
self.fullyConnected = nn.Sequential(
nn.Flatten(),
nn.Linear(in_features=7 * 7 * 64, out_features=128),
nn.ReLU(),
nn.Linear(in_features=128, out_features=10),
)
def forward(self, img):
"""
定义CNN的前向传播路径。
参数：
img -- 输入图像张量
返回值：
output -- 经过CNN模型处理后的输出张量
"""
output = self.conv1(img) # 应用第一个卷积层
output = self.conv2(output) # 应用第二个卷积层
output = self.conv3(output) # 应用第三个卷积层
output = self.fullyConnected(output) # 应用全连接层
return output
def get_device():
if torch.cuda.is_available():
train_device = torch.device('cuda')
else:
train_device = torch.device('cpu')
return train_device
def get_data_loader(dat_path, bat_size, trans, to_train=False):
"""
获取MNIST数据集的数据加载器。
参数:
- dat_path: 数据集的路径。
- bat_size: 批量大小。
- trans: 数据转换器，用于数据预处理。
- to_train: 是否为训练模式，默认为False，即为测试模式。
返回值:
- dat_set: MNIST数据集。
- dat_loader: 数据加载器，用于遍历数据集。
"""
# 加载MNIST数据集，根据to_train参数决定是训练集还是测试集
dat_set = torchvision.datasets.MNIST(root=dat_path, train=to_train, transform=trans, download=True)
# 根据是否为训练模式，配置不同的数据加载器参数
if to_train is True:
dat_loader = torch.utils.data.DataLoader(dat_set, batch_size=bat_size, shuffle=True)
else:
dat_loader = torch.utils.data.DataLoader(dat_set, batch_size=bat_size)
return dat_set, dat_loader
def show_part_of_image(dat_loader, row, col):
"""
展示数据加载器中部分图像的网格视图。
参数:
- dat_loader: 数据加载器，应包含图像和对应的标签。
- row: 每行展示的图像数量。
- col: 每列展示的图像数量。
返回值:
- 无
"""
# 从数据加载器中获取第一个样本
iteration = enumerate(dat_loader)
idx, (exam_img, exam_label) = next(iteration)
# 创建一个图形窗口
fig = plt.figure(num=1)
# 在图形窗口中布局子图
for i in range(row * col):
plt.subplot(row, col, i + 1)
plt.tight_layout()
# 显示图像
plt.imshow(exam_img[i][0], cmap='gray', interpolation='none')
plt.title('Number: {}'.format(exam_label[i])) # 设置图像标题为对应的标签
plt.xticks([]) # 移除x轴刻度
plt.yticks([]) # 移除y轴刻度
plt.show() # 显示图像
def train(network, dat_loader, device, epos, loss_function, optimizer):
"""
对给定的网络进行训练。
参数:
- network: 待训练的网络模型。
- dat_loader: 数据加载器，用于批量加载训练数据。
- device: 指定训练过程使用的设备（如CPU或GPU）。
- epos: 训练的总轮数。
- loss_function: 指定的损失函数。
- optimizer: 优化器，用于更新网络参数。
返回值:
- 训练后的网络模型。
"""
for epoch in range(1, epos + 1):
network.train(mode=True) # 将网络设置为训练模式
for idx, (train_img, train_label) in enumerate(dat_loader):
train_img = train_img.to(device) # 将训练图像数据移动到指定设备
train_label = train_label.to(device) # 将训练标签移动到指定设备
outputs = network(train_img) # 通过网络计算输出
optimizer.zero_grad() # 清除之前的梯度
loss = loss_function(outputs, train_label) # 计算损失
loss.backward() # 反向传播计算梯度
optimizer.step() # 使用优化器更新网络参数
# 每隔100个批次打印一次训练状态
if idx % 100 == 0:
cnt = idx * len(train_img) + (epoch - 1) * len(dat_loader.dataset)
print('epoch: {}, [{}/{}({:.0f}%)], loss: {:.6f}'.format(epoch,
idx * len(train_img),
len(dat_loader.dataset),
(100 * cnt) / (
len(dat_loader.dataset) * epos),
loss.item()))
print('------------------------------------------------')
print('Training ended.')
return network
def test(network, dat_loader, device, loss_function):
"""
测试给定的网络模型。
参数:
- network: 训练好的神经网络模型。
- dat_loader: 测试数据集的数据加载器。
- device: 指定运行设备，如"cpu"或"cuda:0"。
- loss_function: 用于计算损失的函数。
返回值:
- 无返回值，但会打印测试集的损失、正确率等信息。
"""
test_loss_avg, correct, total = 0, 0, 0
test_loss = []
network.train(mode=False) # 将网络设置为评估模式
with torch.no_grad(): # 禁止计算梯度，以减少内存消耗
for idx, (test_img, test_label) in enumerate(dat_loader):
test_img = test_img.to(device) # 将测试图像数据移动到指定设备
test_label = test_label.to(device) # 将测试标签移动到指定设备
total += test_label.size(0) # 统计测试样本总数
outputs = network(test_img) # 通过网络获得输出结果
loss = loss_function(outputs, test_label) # 计算损失
test_loss.append(loss.item()) # 记录当前样本的损失
predictions = torch.argmax(outputs, dim=1) # 获取预测标签
correct += torch.sum(predictions == test_label) # 统计正确预测的数量
test_loss_avg = np.average(test_loss) # 计算测试损失的平均值
# 打印测试结果: 总数、正确数、准确率和平均损失
print('Total: {}, Correct: {}, Accuracy: {:.2f}%, AverageLoss: {:.6f}'.format(total, correct,
correct / total * 100,
test_loss_avg))
def show_part_of_test_result(network, dat_loader, row, col):
"""
展示测试结果的一部分。
参数:
- network: 训练好的神经网络模型。
- dat_loader: 测试数据集的加载器。
- row: 展示图片的行数。
- col: 展示图片的列数。
返回值:
- 无。
"""
# 初始化测试数据的迭代器
iteration = enumerate(dat_loader)
# 获取第一批测试数据
idx, (exam_img, exam_label) = next(iteration)
# 在不计算梯度的情况下通过模型预测输出
with torch.no_grad():
outputs = network(exam_img)
# 创建一个图形用于展示图片
fig = plt.figure()
# 遍历要展示的图片并展示在图形上
for i in range(row * col):
plt.subplot(row, col, i + 1) # 创建子图
plt.tight_layout() # 调整子图间的布局以避免重叠
plt.imshow(exam_img[i][0], cmap='gray', interpolation='none') # 展示图片
plt.title('Number: {}, Prediction: {}'.format(
exam_label[i], outputs.data.max(1, keepdim=True)[1][i].item() # 标题包括真实标签和预测标签
))
plt.xticks([]) # 移除x轴标签
plt.yticks([]) # 移除y轴标签
plt.show() # 展示图形
if __name__ == "__main__":
# 设置批次大小和训练轮数
batch_size, epochs = 64, 12
# 定义数据转换器，包括转换为张量和标准化
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.1307], std=[0.3081])])
# 获取训练使用的设备（如GPU或CPU）
my_device = get_device()
# 数据集路径
path = './data'
# 加载训练数据集和数据加载器
_, train_data_loader = get_data_loader(path, batch_size, transform, True)
print('Training data loaded.')
# 展示训练数据集中的一部分图像
show_part_of_image(train_data_loader, 3, 3)
# 加载测试数据集和数据加载器
_, test_data_loader = get_data_loader(path, batch_size, transform)
print('Testing data loaded.')
# 定义卷积神经网络
cnn = CNN()
# 定义损失函数和优化器
loss_func = nn.CrossEntropyLoss()
optim = torch.optim.Adam(cnn.parameters(), lr=0.01)
# 训练CNN模型
cnn = train(cnn, train_data_loader, my_device, epochs, loss_func, optim)
# 在测试数据集上测试训练好的模型
test(cnn, test_data_loader, my_device, loss_func)
# 展示测试结果的一部分
show_part_of_test_result(cnn, test_data_loader, 5, 2)
# 保存训练好的CNN模型
torch.save(cnn, './cnn2.pth')
# 转化为onnx模型保存
torch.onnx.export(cnn, torch.randn(1, 1, 28, 28), "./cnn2.onnx")
八、开始炼丹
我们直接运行程序就可以开始炼丹了，注意最后两步我们将模型保存为了*.pth以及转换的ONNX模型。
训练开始时我们可以看到右边出现了我们数据集的照片，还有他们的标签，可以说图片还是有些抽象的，因为这是外国人的数据集，所以我们看起来还是有点不适应的，因此我们检测的时候也需要尽量写的跟他们差不多才能达到较高的准确率。
我们可以修改batch_size, epochs来修改网络的训练时间，选个适当的值就行了，batch_size太大会爆内存epochs太大会训练很久而且可能过拟合，所以差不多十几次就够了：
等待训练结束就可以看到我们的模型准确率了，96.83emm，应该是有点过拟合了，之前训练10次的时候还有97的准确率，所以参数的设置也是训练的一个重点哦：
最终我们得到的模型就是下面这两个东西啦：
可以看到这两个文件的大小只有1.6MB，属于是非常小的模型了，这是因为我们网络创建的时候比较小，所以需要的权重参数也比较少，更适合部署在嵌入式这种存储空间少且性能较弱的移动设备上。
九、测试模型
模型都出来了，不用用咋知道好坏呢，随便写个模型加载和测试的代码：
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
import torch.nn as nn
import matplotlib.pyplot as plt
import cv2
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
)
self.fullyConnected = nn.Sequential(
nn.Flatten(),
nn.Linear(in_features=7 * 7 * 64, out_features=128),
nn.ReLU(),
nn.Linear(in_features=128, out_features=10),
)
def forward(self, input):
output = self.conv1(input)
output = self.conv2(output)
output = self.conv3(output)
output = self.fullyConnected(output)
return output
def preprocess_and_highlight_text(img_pil):
# 将PIL Image转换为OpenCV格式（numpy.ndarray）
img_cv = np.array(img_pil.convert('RGB'))
# 预处理（可选）
img_cv = cv2.fastNlMeansDenoising(img_cv, None, 10, 7, 21)
img_cv = cv2.medianBlur(img_cv, 3)
# 二值化
_, img_bin = cv2.threshold(cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY), 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
# 反转颜色
img_highlighted = cv2.bitwise_not(img_bin)
# 将处理后的OpenCV图像转换回PIL Image格式
img_highlighted_pil = Image.fromarray(img_highlighted, mode='L')
return img_highlighted_pil
model = torch.load('./cnn2.pth')
model.eval()
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.1307], std=[0.3081])])
unloader = transforms.ToPILImage()
for k in range(10):
infile = './data/raw/' + '{}.jpg'.format(k)
img = Image.open(infile)
# 突出文字效果
img_highlighted = preprocess_and_highlight_text(img)
# 将图片设置为28*28大小
img = img.resize((28, 28))
img = img.convert('L')
img_array = np.array(img)
# 像素反转
for i in range(28):
for j in range(28):
img_array[i, j] = 255 - img_array[i, j]
# print(img_array)
img = Image.fromarray(img_array)
# img.show()
img = transform(img)
img = torch.unsqueeze(img, 0)
output = model(img)
pred = torch.argmax(output, dim=1)
image = torch.squeeze(img, 0)
image = unloader(image)
plt.subplot(5, 2, k + 1)
plt.tight_layout()
plt.imshow(image, cmap='gray', interpolation='none')
plt.title("Number: {}, Prediction: {}".format(k, pred.item()))
plt.xticks([])
plt.yticks([])
plt.show()
preprocess_and_highlight_text这个函数是用来对图像进行预处理的，不然我们的图像一般都噪声比较多还有其他因素干扰，我的纸偏黄，所以给他进行二值化编程黑白的后再进行反转就跟数据集差不多的，就可以交给模型去识别了，不然你就会得到一堆你都看不懂的数字图片跟别说让模型去看了，测试结果：
效果也还不错了，只有9错了，也还可以。
希望大家都能练出属于自己的丹，给自己折腾小玩意用。
最后偷偷的放一下没有预处理的图片是什么后果：
这告诉我们什么，对于测试数也要进行预处理，还要进行正确的预处理，不然就那阴影一块一块的图片别说模型认不出了，我都不一定认得出来。我一开始还在想模型怎么准确率这么低，看了一眼才发现输入的都是什么妖魔鬼怪QAQ……
|
提升卡
变色卡
千斤顶
1/8 
京公网安备 11010802033920号
Copyright © 2005-2025 EEWORLD.com.cn, Inc. All rights reserved
