手写数字识别(Pytorch)

 一、前言

        现在也算是一名正式的“研0”了，最近也在学一些组里要求掌握的东西，大致的学习方向就是深度学习，机器学习方面的，最近这一阶段在学习pytorch，现在也算是掌握了一些pytorch基础，对于模型的训练套路有了一定的理解，但是不动手实操一下感觉缺点东西，于是就拿这个深度学习入门项目——手写数字识别，来练练手。

二、项目结构

        项目结构设计较为简单，总共分为三个.py文件，它们的作用分别是定义网络模型，定义功能训练模型和在测试集上测试模型准确度，此外整个项目还包括一些其他的文件夹用来存放比如说数据集或者是训练好的模型这样的数据文件，如图：

• dataset：存放下载好的训练数据集或者测试数据集
• logs：用来存放日志文件，用于tensorboard绘图使用
• save_models：存放在验证数据集上效果好的模型
• src：为代码撰写文件夹    

三、功能描述

1.网络模型

        网络模型的设计较为简单，遵循经典的卷积网络模型配置，一个卷积层+一个激活函数+一个池化层+展平+全连接层

import torch
from torch import nn


class CNNNet(nn.Module):
    def __init__(self):
        super(CNNNet, self).__init__()
        self.covn1 = nn.Sequential(
            nn.Conv2d(1, 10, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(10, 20, kernel_size=4),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(500, 50),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(50, 10),
            nn.LogSoftmax(dim=1)
        )
    def forward(self, x):
        x = self.covn1(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

        输入这个模型是一个四维矩阵[256, 1, 28, 28]，256张图片，每个图片一个通道，大小是28*28，经过网络模型后会变为二维矩阵[256, 10]，每一行代表一个图片，每一列表示一个特征，在这个具体问题下，每一列表示为某个数的概率。

1.数据集处理

        首先是从网络中下载MNIST数据集，这里使用torchvision的API进行下载，由于数据集不大，速度相对较快，对于大的数据集推荐使用迅雷进行下载，然后再将解压好的数据包导入到dataset文件夹中。

import math

import numpy as np
import torch.cuda
import torchvision
from torch import nn
from torch.utils.data import Dataset, DataLoader, random_split
from torch.utils.tensorboard import SummaryWriter

from tqdm import tqdm


dataset_transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.1307,), (0.3081,))
])

train_set = torchvision.datasets.MNIST(root='../dataset', train=True, transform=dataset_transform, download=True)
test_set = torchvision.datasets.MNIST(root='../dataset', train=False, transform=dataset_transform, download=True)

         对于读取的数据集可以同时进行初始化操作，将数据转化为张量并且进行归一化处理，此均值和方差常用于MINST数据集，训练模型时可以加速收敛并且可以消除图像像素数据量纲的影响。

2.超参数设置

        在深度学习中有很多超参数需要自己手动设定，我们可以将需要设定的超参数全部整合成一个字典，这样在使用超参数时通过字典可以直接读取，当修改某一个超参数的时可以直接修改字典内容，方便管理，同样方便修改。

device = "cuda" if torch.cuda.is_available() else "cpu"
config = {
    'seed': 5201314, #随机数种子
    'valid_ratio': 0.2,  # validation_size = train_size * valid_ratio
    'n_epochs': 10, 
    'batch_size': 256,
    'learning_rate': 1e-2,
    'early_stop': 400,  # 模型训练提前结束阈值
    'save_path': '../save_ models/network.pth'
}

        此外，为了方便调试，可以通过随机数种子来使得每次运行得到的随机数都是同一个数，并且取消模型随机优化，保证相同的输入得到相同的输出，这对于调试程序很重要。

def same_seed(seed):
    torch.backends.cudnn.deterministic = True # 禁用可能引入非确定性元素的某些算法来强制CuDNN中的确定性行为
    torch.backends.cudnn.benchmark = False # 用运行时的 CuDNN 自动调整，并使用一组确定的算法，这有助于确保在不同运行之间获得一致的性能
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

 3.划分训练数据集和验证数据集

        在训练结束后可以使用验证数据集对训练的模型进行验证，观察在验证数据集上的损失值来选择较为优秀的模型，以供测试数据集使用。可以使用K折交叉验证来评估模型的性能或者是调整模型参数。

def train_valid_split(data_set, valid_ratio, seed):
    valid_set_size = int(valid_ratio * len(data_set))
    train_set_size = len(data_set) - valid_set_size
    train_set, valid_set = random_split(data_set, [train_set_size, valid_set_size], generator=torch.Generator().manual_seed(seed))
    return train_set, valid_set

        经过randomsplit函数划分后，train_set和valid_set的数据类型还是dataset类型，可以作为DataLoader类的参数。

4.模型训练函数

        模型训练函数是最重要的功能，包括损失函数的设定，优化器的选择等，对于手写数字识别这种经典的分类问题可以使用交叉熵损失函数作为损失函数，优化器选择Adam，Adam可以动态调整学习率，降低出现局部最小值或者是鞍点的可能。

def trainer(train_loader, valid_loader, model, config, device):
    criterion = nn.CrossEntropyLoss(reduction='mean')
    optimizer = torch.optim.Adam(model.parameters(), lr=config['learning_rate'])
    writer = SummaryWriter("../logs") # 日志保存路径
    n_epochs, best_loss, step, early_stop_count, model_version = config['n_epochs'], math.inf, 0, 0, 0
    for epoch in range(n_epochs):
        model.train() # 训练
        loss_record = []

        train_pbar = tqdm(train_loader, position=0, leave=True) # 用进度条来显示模型训练过程
        for inputs, labels in train_pbar:
            optimizer.zero_grad()
            inputs, labels = inputs.to(device), labels.to(device)
            pred = model(inputs)
            loss = criterion(pred, labels)
            loss.backward()
            optimizer.step()
            step += 1
            loss_record.append(loss.detach().item()) # 记录损失函数值，值于梯度分离
            train_pbar.set_description(f"Epoch[{epoch+1}/{n_epochs}]")
            train_pbar.set_postfix({"loss" : loss.detach().item()})

        mean_train_loss = sum(loss_record)/len(loss_record) # 计算平均损失
        writer.add_scalar("train", mean_train_loss, step) # 使用tensorboard绘制每个epoch下平均训练损失值曲线

        model.eval() # 验证
        loss_record = []
        for x, y in valid_loader:
            x, y = x.to(device), y.to(device)
            with torch.no_grad():
                pred = model(x)
                loss = criterion(pred, y)

            loss_record.append(loss.item())

        mean_valid_loss = sum(loss_record)/len(loss_record)
        print(f"Epoch[{epoch + 1}/{n_epochs}]: Train loss: {mean_train_loss:.4f}, Valid loss:{mean_valid_loss:.4f}")
        writer.add_scalar("valid", mean_valid_loss, step) # 使用tensorboard绘制每个epoch下平均验证损失值曲线

        # 保存在验证数据集上性能最好的模型
        if mean_valid_loss < best_loss: 
            best_loss = mean_valid_loss
            torch.save(model.state_dict(), f"../save_models/model_version_{model_version + 1}.pth")
            print(f"Saving model with loss {best_loss:.3f}")
            early_stop_count = 0
        else:
            early_stop_count += 1

        # 模型长时间没有得到优化时，提前结束训练
        if early_stop_count >= config["early_stop"]:
            print("\nModel is not improving, so we halt the training session")
            writer.close()
            return

        model_version += 1

5.开始训练 

same_seed(config['seed'])
train_data, valid_data = train_valid_split(train_set, config['valid_ratio'], config['seed'])
# print(type(train_data), '\n', len(train_data))
# print(type(valid_data), '\n', len(valid_data))
train_loader = DataLoader(train_data, batch_size=config['batch_size'], shuffle=True, pin_memory=True)
valid_loader = DataLoader(valid_data, batch_size=config['batch_size'], shuffle=True, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=config['batch_size'], shuffle=True,pin_memory=True)

model = CNNNet().to(device)
# print(model)
trainer(train_loader, valid_loader, model, config, device)

训练结果如图：

​​​​​​​​​​​​​​

        比较损失函数，第九次的训练效果更好，保存模型。

6.模型测试

        读取保存好的模型，在测试数据集上运行，检验正确率

import numpy as np
import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader

from src.My_module import CNNNet

model = CNNNet()
model.load_state_dict(torch.load("../save_models/model_version_9.pth"))

test_loss_avg = 0

def test(model, test_loader, device='cpu'):
    correct = 0
    total = 0
    test_loss = []
    criterion = nn.CrossEntropyLoss(reduction='mean')
    with torch.no_grad():
        for train_idx, (inputs, labels) in enumerate(test_loader, 0):

            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            test_loss.append(loss.item())
            index, value = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += int((value==labels).sum())

        test_loss_avg = np.average(test_loss)
        print('Total: {}, Correct: {}, Accuracy: {:.2f}%, AverageLoss: {:.6f}'.format(total, correct, (correct/total*100), test_loss_avg))


dataset_transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.1307,), (0.3081,))
])

config = {
    'seed': 5201314,  # Your seed number, you can pick your lucky number. :)
    'valid_ratio': 0.2,  # validation_size = train_size * valid_ratio
    'n_epochs': 10,  # Number of epochs.
    'batch_size': 256,
    'learning_rate': 1e-2,
    'early_stop': 400,  # If model has not improved for this many consecutive epochs, stop training.
    'save_path': '../save_ models/network.pth'  # Your model will be saved here.
}

test_set = torchvision.datasets.MNIST(root='../dataset', train=False, transform=dataset_transform, download=True)
test_loader = DataLoader(test_set, batch_size=config['batch_size'], shuffle=True,pin_memory=True)

test(model, test_loader)

Total: 10000, Correct: 9260, Accuracy: 92.60%, AverageLoss: 0.250599

        训练次数还是少，有时间的小伙伴可以增大config中n_epoch，增加训练的次数，有可能会得到更好的结果 。

四、总结

        在pytorch基础阶段，这个小项目覆盖的内容还是比较全面的，包括，数据集读取加载，数据处理，网络模型设计，模型训练等不少东西，拿来练手熟悉套路是一个很好的选择，当然也是借鉴了其他大佬的代码（整体结构的设计思路参考李宏毅老师机器学习作业一），外加自己的一些理解，完成了这次小项目。相应的参考我会放在下面，小伙伴们可以对比着选择性的查看。

参考：http://t.csdnimg.cn/AViuM