pytorch手写数字识别（CNN）

---

一、前言

1. 数据集是 MNIST手写数字 数据集，自取 https://download.csdn.net/download/weixin_43721000/87708518
2. 网络结构为两层卷积两层线性层

---

二、实现方法

1.导包

import numpy as np 
import pandas as pd 

from matplotlib import pyplot as plt

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

2. 使用gpu

# 使用 GPU ----------------------------------------------
print(torch.cuda.is_available())
device = torch.device("cuda")
print(device)

3. 定义数据读取方法

# 定义数据读取类 ------------------------------------------
class MNISTDataset(Dataset):
    def __init__(self, data_type):
        self.data_type = data_type
        
        if self.data_type not in ['train', 'test']: raise Exception('数据集类型有误!')
        
        # 指定数据集路径
        if self.data_type == 'train':
            self.data_path = "/kaggle/input/digit-recognizer/train.csv"
        elif self.data_type == 'test':
            self.data_path = "/kaggle/input/digit-recognizer/test.csv"
        
        # 读取数据集
        df = pd.read_csv(self.data_path)
        
        # 标签和样本赋值
        if self.data_type == 'train':
            # 训练集标签与图片分割
            self.label = df['label'].to_numpy().reshape(-1, 1)
            self.image = df.drop(columns='label').to_numpy().reshape(-1, 1, 28, 28) / 255
        elif self.data_type == 'test':
            # 测试集没有标签，直接加载
            self.image = df.to_numpy().reshape(-1, 1, 28, 28) / 255

            
    def __len__(self):
        return len(self.image)
    
    
    def __getitem__(self, index):
        # 样本转为tensor
        image = torch.FloatTensor(self.image[index])
        
        # 标签转为模型输出格式（仅训练集有标签）
        if self.data_type == 'train':
            label = torch.zeros(10, dtype=torch.float)
            label[int(self.label[index])] = 1.0
            return image, label
        else:
            return image

# 训练集    
train_datasets = MNISTDataset("train")
train_dataloader = DataLoader(train_datasets, batch_size=100, shuffle=True)

# 测试集
test_datasets = MNISTDataset("test")
test_dataloader = DataLoader(test_datasets, batch_size=100, shuffle=False)

4.定义模型、损失函数、优化器

# 定义模型结构 -----------------------------------------
# 两个卷积加两个线性层的网络
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.seq1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),  # 32个卷积核卷积，卷积后图片数量变为32，大小不变 (28-3+1)/1 + 1*2 = 28 
            nn.ELU(inplace=True),                                  # 激活函数
            nn.MaxPool2d(kernel_size=2, stride=2)                  # 宽高为2步长为2进行最大池化，图片宽高减半，变为 14
        )
        
        self.seq2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1), # 64个卷积核卷积，卷积后图片数量变为64，大小不变 (14-3+1)/1 + 1*2 = 14 
            nn.ELU(inplace=True),                                  # 激活函数
            nn.MaxPool2d(kernel_size=2, stride=2)                  # 宽高为2步长为2进行最大池化，图片宽高减半，变为 7
        )
        
        self.seq3 = nn.Sequential(
            nn.Linear(7*7*64, 32, bias=True),                      # 64张图片像素点拉直输入全连接层
            nn.ELU(inplace=True),                                  # 激活函数
            nn.Linear(32, 10, bias=True),                          # 经过全连接输出10分类
            nn.ELU()                                               # 激活函数
        )
        
    def forward(self, x):
        x = self.seq1(x)
        x = self.seq2(x)
        x = x.view(x.size(0), -1)
        x = self.seq3(x)
        return x


# 初始化模型
model = SimpleCNN()
model = model.to(device)

# 定义损失函数（交叉熵）
criterion = nn.CrossEntropyLoss().to(device)

# 定义优化器（adam）
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

5.训练

# 训练 ----------------------------------------------------
# 记录每代平均损失值的数组
history_loss = []
# 迭代100次
EPOCH = 100
for epoch in range(EPOCH):
    loss_tmp = []
    for i, (image, label) in enumerate(train_dataloader):
        # 拿出一个批次的数据
        image, label = image.to(device), label.to(device)
        # 送入网络
        y_pred = model(image)
        # 计算损失
        loss = criterion(y_pred, label)
        # 梯度清零
        optimizer.zero_grad()
        # 反向传播
        loss.backward()
        # 梯度更新
        optimizer.step()
        # 记录损失值
        loss_tmp.append(loss.item())
        
    # 计算当前 EPOCH 的平均损失
    history_loss.append(np.mean(loss_tmp))
    print(f"\rEPOCH: [{epoch+1}/{EPOCH}], LOSS: {history_loss[-1]:1.6f}", end="")

6.绘制损失值图像

# 绘制每代平均损失值图像
plt.plot(history_loss, label="Cross Entropy Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()

7.预测

# 对所有测试集进行预测
result = []
for i, image in enumerate(test_dataloader):
    # 图片转到 GPU 上
    image = image.to(device)
    # 预测
    y_pred = model(image)
    label_pred = torch.argmax(y_pred, dim=1)
    # 预测结果转回cpu并加入数组
    result += label_pred.cpu().tolist()

8.绘制测试集图像与预测结果的对照图

# 绘制前 30 项预测结果
height, width = 5, 6
fig, axes = plt.subplots(height, width, figsize=(10, 10))

for i in range(height):
    for j in range(width):
        idx = height * i + j
        axes[i][j].imshow(test_datasets[idx][0], cmap='gray')
        axes[i][j].set_title(result[idx])
        axes[i][j].axis('off')

plt.show()