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))))
554
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
