CNN识别MNIST手写字

开源地址：https://github.com/MYJOKERML/CNN

学了这么久，终于可以自己搭建一个CNN了，还记得大一时看了好多次这样的代码，然后全都无功而返。现在会看去年的自己，哈，这小家伙居然连个python代码都看不懂，连个class都不会。也许再等一年后再次回头时我可以笑着说，哈，一年前的自己居然看一个小时论文要翻半个小时字典，自己的进步真的肉眼可见啊，也得益于日益进步的科技，在copilot的协助下，现在普通的代码已经主要以理清逻辑为主了，人类的生产力又再一次有了质的飞跃，真的了不起。

话不多说，直接上代码记录自己的学习过程。

导入相关包

import torch
from torch import nn, optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader, random_split

from matplotlib import pyplot as plt
from PIL import Image

划分数据集并实现数据增强

这里数据增强我才用了常见的数据增强办法，随机旋转图片和平移图片，以及颜色变换。我注释掉了随机水平翻转，因为注意到很多数字水平翻转后并没有任何意义。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 定义数据预处理转换，包括数据增强操作
transform = transforms.Compose([
   transforms.RandomRotation(degrees=15),  # 随机旋转
   transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),  # 随机平移
   transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 随机颜色变换
   # transforms.RandomHorizontalFlip(),  # 随机水平翻转
   transforms.ToTensor(),
])

dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)

# 定义训练集和测试集的比例
train_ration = 0.8
test_ration = 1 - train_ration

# 计算训练集和测试集的数量
train_size = int(train_ration * len(dataset))
test_size = len(dataset) - train_size

# 按照计算的数量随机划分训练集和测试集
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

加载数据集并查看第一张图片

# 加载数据并查看数据
# from einops import rearrange

BATCH_SIZE = 64
train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, shuffle=False)

for i, (images, labels) in enumerate(train_loader):
   print(i, images.shape, labels.shape)
   plt.imshow(images[0][0], cmap='gray')
   plt.title(labels[0].item())
   break

搭建卷积神经网络

卷积核一般是3*3的大小，于是令padding=1就可以保持原图像大小不变。

最开始只搭了两层卷积层，然后2层全连接层，事实证明这样非常容易过拟合，通常第一轮epoch结束后就有90%几的正确率，5轮之后会有97%的正确率，显然过拟合了，我自己手写了一张 8 的图片，训练了好多次一直都是判断为 0 。。。

后来我想明白了，结果不理想的原因是网络深度太浅，根本提取不到足够大小范围的图片，于是后面搭了4层卷积层，3层全连接层，效果确实变好了。

最加深网络时特征提取多了（即 out_channels 设大了），于是模型很难训练，正确率一直在70%几上不去，想了想应该是因为本来图片尺寸就小28*28，又是灰度图，没这么多特征，所以最后输出通道数就比较小，成功改善了结果。最后训练出来正确率在94%左右。

本来一般是有加池化层的，但现在算力足够了，也就没加池化层，完全可以直接算。

class CNN(nn.Module):
   def __init__(self):
       super(CNN, self).__init__()
       # 卷积层提取图片特征
       self.conv1 = nn.Sequential(
           nn.Conv2d(1, 16, 3, 1, 1),
           nn.ReLU()
       )
       self.conv2 = nn.Sequential(
           nn.Conv2d(16, 16, 3, 1, 1),
           nn.ReLU()
       )

       # 添加更多的卷积层
       self.conv3 = nn.Sequential(
           nn.Conv2d(16, 32, 3, 1, 1),
           nn.ReLU()
       )
       self.conv4 = nn.Sequential(
           nn.Conv2d(32, 32, 3, 1, 1),
           nn.ReLU()
       )

       # 添加池化层
       # self.pool = nn.MaxPool2d(2, 2)

       self.fc_input_size = 32 * 28 * 28

       self.fc1 = nn.Sequential(nn.Linear(self.fc_input_size, 256), nn.Dropout(0.2), nn.ReLU())
       self.fc2 = nn.Sequential(nn.Linear(256, 128), nn.Dropout(0.2), nn.ReLU())
       self.fc3 = nn.Sequential(nn.Linear(128, 10), nn.Softmax(dim=1))

   def forward(self, x):
       x = self.conv1(x)
       # x = self.pool(x)
       x = self.conv2(x)
       # x = self.pool(x)
       x = self.conv3(x)
       # x = self.pool(x)
       x = self.conv4(x)
       # x = self.pool(x)

       x = x.view(x.shape[0], -1)

       x = self.fc1(x)
       x = self.fc2(x)
       x = self.fc3(x)

       return x

定义损失函数和优化器

# 定义损失函数和优化器
model = CNN()
model = model.to(device)
criterion = nn.CrossEntropyLoss()
lr=0.001 # 学习率
optimizer = optim.Adam(model.parameters(), lr)

训练模型和测试模型

from tqdm import tqdm

best_loss = float('inf')  # 初始化为正无穷大，确保第一个损失值一定会小于它

# 训练模型
def train():
   model.train()
   total_loss = 0 # 用于计算平均损失
   num_batches = len(train_loader) # 用于记录训练的batch数目

   with tqdm(total=num_batches, desc='Training', unit='batch') as pbar:
       for i, data in enumerate(train_loader):
           inputs, labels = data
           inputs, labels = inputs.to(device), labels.to(device)
           optimizer.zero_grad() # 梯度清零
           outputs = model(inputs)
           loss = criterion(outputs, labels)
           loss.backward() # 反向传播
           optimizer.step() # 更新参数

           total_loss += loss.item()
           current_loss = loss.item()

           if i % 100 == 0:
               # print('Train Step: {}\tLoss: {:.3f}'.format(i, loss.item()))
               avg_loss = total_loss / (i + 1)
               pbar.set_postfix({'cur_loss': '{:.3f}'.format(loss), 'avg_loss':'{:.3f}'.format(avg_loss)})
               pbar.update(100) # 每处理100个batch更新一次tqdm
               # 判断当前损失是否比最佳损失小，如果是则保存模型状态
           if current_loss < best_loss:
               best_loss = current_loss
               torch.save(model.state_dict(), 'best_model.pth')  # 保存模型

   print('Best loss:', best_loss)            

# 测试模型
def test():
   model.eval()
   correct = 0
   total = 0
   with torch.no_grad(): # 测试过程中不需要计算梯度
       with tqdm(total=len(test_loader), desc="Testing", unit="batch") as pbar:
           for data in test_loader:
               inputs, labels = data
               inputs, labels = inputs.to(device), labels.to(device)
               outputs = model(inputs)
               _, predicted = torch.max(outputs.data, dim=1)
               total += labels.size(0)
               # print(predicted.shape, labels.shape)
               correct += (predicted == labels).sum().item()
               pbar.update(1)# 每处理1个batch更新一次tqdm

   print('Accuracy on test set: {:.3f}'.format(correct / total))

print('Before training:', end=' ')
test()

保存和加载模型

上面已经保存了loss最低的模型，但可能过拟合，因此再保存一个最后一次运行的模型，加载loss最低的模型进行测试。

# 保存最后一次模型

torch.save(model.state_dict(), os.path.join(save_path, 'model.pth'))

# 加载模型
model = CNN()
model.load_state_dict(torch.load(os.path.join(save_path, 'best_model.pth')))
model = model.to(device)

# 预测
def predict(img):
   model.eval()
   img = img.to(device)
   with torch.no_grad():
       output = model(img)
       # print(output)
       _, predicted = torch.max(output.data, dim=1)
       return predicted.item()

import random
# 随机选取一张图片进行预测
index = random.randint(0, len(test_dataset))
img = test_dataset[index][0].unsqueeze(0)
label = test_dataset[index][1]
plt.imshow(img[0][0], cmap='gray')
plt.title('label: {}'.format(label))
plt.show()
print('predict: {}'.format(predict(img)))

自己手写了一张数字 8 的图片进行测试。

先对图片进行缩放，再转换为灰度图，再转换成张量之后输入模型进行预测。之前一直预测为0，但更改模型深度后就可以较好地预测出该图片为8。

from PIL import Image

img = Image.open('./8.jpg')

transform = transforms.Compose([
   transforms.Resize((28, 28)),
   transforms.Grayscale(),
   transforms.ToTensor(),
])

print("Original size: ", img.size)
img = transform(img)
print(img.shape)
plt.title('label: 8')
plt.imshow(img[0], cmap='gray')
plt.show()

print('predict: {}'.format(predict(img.unsqueeze(0))))

Jerry的小屋

CNN识别MNIST手写字

CNN识别MNIST手写字

konachan爬取图片

下一篇文章

Jerry

叨叨几句... NOTHING

取消回复

Jerry的小屋