Pytorch实现卷积神经网络手写数字识别(MNIST)

本文使用pytorch完成一个非常经典的任务——手写数字识别。数据集为LeCun等人于90年代构建的手写数字集MNIST。本文的重点在于对数据的处理，因为torchvision库虽然可以直接加载训练数据，但是这个过程中是看不见数据的，所以本文使用dataset和dataloader库来加载数据，可以从中体会数据处理的过程。

        首先引用本项目需要用到的所有库，并定义本项目的一些超参数。

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torch.optim.lr_scheduler import ExponentialLR
import torch.nn as nn


epochs = 10
batch_size = 128
learning_rate = 0.01

1 数据处理

1.1 数据集结构

        MNIST数据集可以在官方页面下载，整个数据集分为训练/测试图片和训练/测试标签四个文件。这四个文件以二进制的形式存储，对于图片文件来说，它的前16位是图片集的各种信息，后面按顺序存储着图片的像素信息；对于标签文件来说，它的前8位是标签集的各种信息，后面存储着标签值。

1.2 加载训练数据

        下载并解压四个文件，即可开始读取数据。根据1.1中的介绍，我们对数据集处理后读入numpy数组。

#数据集路径
images_train = 'D:\\datasets\\MNIST\\train-images.idx3-ubyte'
images_test = 'D:\\datasets\\MNIST\\t10k-images.idx3-ubyte'
labels_train = 'D:\\datasets\\MNIST\\train-labels.idx1-ubyte'
labels_test = 'D:\\datasets\\MNIST\\t10k-labels.idx1-ubyte'


#读取数据到numpy数组
def read_images(dir):
    with open(dir, 'rb') as f:
        f.read(16)
        tmp = f.read()
        images = np.frombuffer(tmp, dtype=np.uint8).astype("float32").reshape(int(len(tmp)/(28*28)), 1, 28, 28)
    return images


def read_labels(dir):
    with open(dir, 'rb') as f:
        f.read(8)
        tmp = f.read()
        labels = np.frombuffer(tmp, dtype=np.uint8).astype("int64")
    return labels

        经过函数read_images的处理，图片集将会被处理为 (C, 1, H, W) 的numpy数组。C表示图片的个数；对于每张图片，它的尺寸是 (1, H, W) ，H、W分别为图片高和宽。将图片集处理成 (C, 1, H, W) 而不是 (C, H, W) 的原因是，将来网络提取的特征是高维的，需要预留一个数组维度。

        经过函数read_labels的处理，标签集将会被处理为一维向量。

        此时，我们分别读取到了图像和标签信息。接下来，我们需要一一建立图片与标签的对应关系，这里可以用pytorch中的Dataset（torch.utils.data.Dataset）来完成。注意，我们在建立数据集时，首先要继承Dataset类，并且根据官方文档的说明，我们还需要重写__getitem__()和__len__()。可以看到，图像和标签的对应关系在__getitem__()中构建好了。

        建立好对应关系后，我们还需要通过DataLoader（torch.utils.data.Dataloader）构建一个迭代器，它可以把整个数据集按照batch_size分成多个batch，在训练时迭代地把这些数据按批次送入网络。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Running on", device, '\n')


class dataset(Dataset):
    def __init__(self, images, labels):
        self.images = torch.tensor(images).to(device)
        self.labels = torch.tensor(labels).to(device)

    def __getitem__(self, item):
        return self.images[item], self.labels[item]

    def __len__(self):
        return self.images.shape[0]


train_data = dataset(read_images(images_train), read_labels(labels_train))
train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, drop_last=False)
test_data = dataset(read_images(images_test), read_labels(labels_test))
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False, drop_last=False)

        这段代码前面两行指定了代码运行的设备。这里应该注意，要使用gpu，我们就需要同时把数据和网络权重都储存在gpu上。

2 网络构建

        网络部分是该任务的重点，这一部分可以自行设计，例如我自己设计的MyNet：

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        #(1, 28, 28) -> (32, 12, 12)
        self.conv1 = nn.Sequential(nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5, stride=1),
                                   nn.BatchNorm2d(32),
                                   nn.ReLU(),
                                   nn.MaxPool2d(kernel_size=2))

        #(32, 12, 12) -> (64, 23, 23)
        self.convtrans = nn.Sequential(nn.ConvTranspose2d(in_channels=32, out_channels=64, kernel_size=2, stride=4),
                                       nn.BatchNorm2d(64),
                                       nn.ReLU(),
                                       nn.MaxPool2d(kernel_size=2))

        #(64, 23, 23) -> (128, 5, 5)
        self.conv2 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=2),
                                   nn.BatchNorm2d(128),
                                   nn.ReLU(),
                                   nn.MaxPool2d(kernel_size=2))

        #(128, 5, 5) -> (3200,)
        self.flatten = nn.Flatten(start_dim=1, end_dim=-1)

        self.linear = nn.Linear(3200, 10)

    def forward(self, input):
        input = self.conv1(input)
        input = self.convtrans(input)
        input = self.conv2(input)
        input = self.flatten(input)
        output = self.linear(input)
        return output

 3 模型加载

        这一部分实例化定义好的网络，并将模型加载到运算设备(cpu或gpu)上。同时定义了优化器、学习率衰减策略和损失函数。

model = MyNet()
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)
scheduler = ExponentialLR(optimizer, gamma=0.7)
loss_calculator = nn.CrossEntropyLoss().to(device)

 4 训练及测试

        本文中训练100个batch就执行一次测试，用来观察训练情况。训练及测试代码如下：

def train(train_dataloader, test_dataloader, optimizer, loss_calculator):
    for epoch in range(epochs):
        for batch, (images, labels) in enumerate(train_dataloader):
            preds = model(images)
            loss = loss_calculator(preds, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if batch % 100 == 0:
                print("Epoch: {: <5}Batch: {: <8}Train_loss: {: <15.6f}".format(epoch, batch, float(loss.data)), end='')
                test(test_dataloader)
        scheduler.step()


def test(test_dataloader):
    match = 0
    total = 0
    with torch.no_grad():
        for imgs, labels in test_dataloader:
            preds = model(imgs)
            preds = torch.max(preds, 1)[1]
            match += int((preds == labels).sum())
            total += int(labels.size()[0])
    accuracy = match/total
    print('Test accuracy: {:.2f}%'.format(accuracy*100))

5 运行代码和模型保存

train(train_dataloader, test_dataloader, optimizer, loss_calculator)
torch.save(model, "MyNet_BS{}_LR{}_ExpLR{}.pth".format(batch_size, learning_rate, 0.7))

模型的保存和加载有很多方式，本文不再赘述。

本文完整代码在这里：https://download.csdn.net/download/diqiudq/85593886，积分不够的话可以把以上所有代码段复制拼接到一个文档运行。