【深度学习】基于华为MindSpore和pytorch的卷积神经网络LeNet5实现MNIST手写识别

1 实验内容简介

1.1 实验目的

（1）熟练掌握卷积、池化概念；

（2）熟练掌握卷积神经网络的基本原理；

（3）熟练掌握各种卷积神经网络框架单元；

（4）熟练掌握经典卷积神经网络模型。

1.2 实验内容及要求

请基于pytorch和mindspore平台，利用MNIST数据集，选择一个典型卷积模型，构建一个自己的卷积模型，以分类的准确度和混淆矩阵为衡量指标，分析两个模型的分类精度。

要求：pytorch可与tensorflow替换，但mindspore为必选平台。（mindspore可以在华为云ModelArts上实现）；自己构建的模型必须与经典模型的网络结构有显著区分，但总体分类准确度需要在97%以上记为合格（不扣分）。

1.3 实验数据集介绍

1.3.1 数据集简介

MNIST数据集（Mixed National Institute of Standards and Technology Database）是一个用来训练各种图像处理系统的二进制图像数据集，广泛应用于机器学习中的训练和测试。MNIST数据集共有70000张图像，其中训练集60000张，测试集10000张。所有图像都是28×28的灰度图像，每张图像包含一个手写数字。

1.3.2 数据集详细信息

（1）数据量

训练集60000张图像，其中30000张来自NIST的Special Database 3，30000张来自NIST的Special Database 1。测试集10000张图像，其中5000张来自NIST的Special Database 3，5000张来自NIST的Special Database 1。

（2）标注情况

每张图像都有标注。

（3）标注类别

共10个类别，每个类别代表0~9之间的一个数字，每张图像只有一个类别。

1.3.3 数据集文件结构

（1）目录结构

·解压前

dataset_compressed/

├── t10k-images-idx3-ubyte.gz #测试集图像压缩包(1648877 bytes)

├── t10k-labels-idx1-ubyte.gz #测试集标签压缩包(4542 bytes)

├── train-images-idx3-ubyte.gz #训练集图像压缩包(9912422 bytes)

└── train-labels-idx1-ubyte.gz #训练集标签压缩包(28881 bytes)

·解压后

dataset_uncompressed/

├── t10k-images-idx3-ubyte #测试集图像数据

├── t10k-labels-idx1-ubyte #测试集标签数据

├── train-images-idx3-ubyte #训练集图像数据

└── train-labels-idx1-ubyte #训练集标签数据

（2）文件结构

MNIST数据集将图像和标签都以矩阵的形式存储于一种称为idx格式的二进制文件中。该数据集的4个二进制文件的存储格式分别如下：

·训练集标签数据 (train-labels-idx1-ubyte)

·训练集图像数据(train-images-idx3-ubyte)

·测试集标签数据(t10k-labels-idx1-ubyte)

·测试集图像数据 (t10k-images-idx3-ubyte)

2 算法原理阐述

2.1 LeNet-5网络结构

LeNet-5由LeCun等人提出于1998年提出，是一种用于手写体字符识别的非常高效的卷积神经网络，出自论文《Gradient-Based Learning Applied to Document Recognition》。LeNet-5的网络结构如下图所示：

整个 LeNet-5 网络总共包括7层（不含输入层），分别是：C1、S2、C3、S4、C5、F6和OUTPUT。输入二维图像（单通道），先经过两次卷积层到池化层，再经过全连接层，最后为输出层。接受输入图像大小为32×32=1024，输出对应10个类别的得分。LeNet-5 中的每一层结构如下：

（1）C1层是卷积层，使用6个5×5的卷积核，得到6组大小为28×28 = 784的特征映射。因此，C1 层的神经元数量为6×784 = 4 704，可训练参数数量为6×25 + 6 = 156，连接数为156×784 = 122304（包括偏置在内，下同）。

（2）S2层为汇聚层，采样窗口为2×2，使用平均汇聚，并使用一个非线性函数。神经元个数为6×14×14 =1176，可训练参数数量为6×(1 + 1) = 12，连接数为6×196×(4 + 1) = 5880。

（3）C3 层为卷积层。LeNet-5 中用一个连接表来定义输入和输出特征映射之间的依赖关系。共使用60个5×5的卷积核，得到16 组大小为10×10的特征映射。如果不使用连接表，则需要96个5×5的卷积核。神经元数量为16 ×100 = 1600，可训练参数数量为(60×25) + 16 = 1516，连接数为100×1516 = 151600。

（4）S4 层是一个汇聚层，采样窗口为2×2，得到16个5×5大小的特征映射，可训练参数数量为16×2 = 32，连接数为16×25×(4 + 1) = 2 000。

（5）C5 层是一个卷积层，使用120×16 = 1 920 个5×5的卷积核，得到120 组大小为1×1的特征映射。C5 层的神经元数量为120，可训练参数数量为1 920×25 + 120 = 48120，连接数为120×(16×25 + 1) = 48120。

（6）F6 层是一个全连接层，有84 个神经元，可训练参数数量为84× (120 +1) = 10164。连接数和可训练参数个数相同，为10164。

（7）输出层：输出层由10 个径向基函数（Radial Basis Function，RBF）组成。

卷积层的每一个输出特征映射都依赖于所有输入特征映射，相当于卷积层的输入和输出特征映射之间是全连接的关系。实际上，这种全连接关系不是必须的。我们可以让每一个输出特征映射都依赖于少数几个输入特征映射。定义一个连接表（Link Table）来描述输入和输出特征映射之间的连接关系。在LeNet-5中，连接表的基本设定如下图所示。C3层的第0-5个特征映射依赖于S2层的特征映射组的每3个连续子集，第6-11个特征映射依赖于S2层的特征映射组的每4个连续子集，第12-14个特征映射依赖于S2层的特征映射的每4个不连续子集，第15个特征映射依赖于S2层的所有特征映射。

3 实验流程及代码实现

3.1 实验平台简介

3.1.1 MindSpore

MindSpore是华为公司自研的最佳匹配昇腾AI处理器算力的全场景深度学习框架，为数据科学家和算法工程师提供设计友好、运行高效的开发体验，推动人工智能软硬件应用生态繁荣发展，目前MindSpore支持在EulerOS、Ubuntu、Windows系统上安装。

3.1.2 pytorch

PyTorch是一个开源的Python机器学习库，基于Torch，用于自然语言处理等应用程序。2017年1月，由Facebook人工智能研究院（FAIR）基于Torch推出了PyTorch。它是一个基于Python的可续计算包，提供两个高级功能：1、具有强大的GPU加速的张量计算（如NumPy）。2、包含自动求导系统的深度神经网络。

3.2 评价指标

3.2.1 混淆矩阵

混淆矩阵（Confusion Matrix）又被称为错误矩阵，通过它可以直观地观察到算法的效果。它的每一列是样本的预测分类，每一行是样本的真实分类（反过来也可以），顾名思义，它反映了分类结果的混淆程度。

·P（Positive）：代表1，表示预测为正样本；

·N（Negative）：代表0，表示预测为负样本；

·T（True）：代表预测正确；

·F（False）：代表预测错误。

下列Positive和Negative表示模型对样本预测的结果是正样本（正例）还是负样本（负例）。True和False表示预测的结果和真实结果是否相同。

·True positives（TP）

预测为1，预测正确，即实际为1；

·False positives（FP）

预测为1，预测错误，即实际为0；

·False negatives（FN）

预测为0，预测错误，即实际为1；

·True negatives（TN）

预测为0，预测正确，即实际为0。

3.2.2 准确率

准确率（Accuracy）衡量的是分类正确的比例。

3.3 实验流程

3.3.1 基于MindSpore的LeNet-5网络

首先导入实验需要的模块。mindspore中context模块用于配置当前执行环境，包括执行模式等特性；vision.c_transforms模块是处理图像增强的高性能模块，用于数据增强图像数据改进训练模型。同时需要设置MindSpore的执行设备和模式。

import os
import mindspore as ms
# 导入mindspore中context模块，用于配置当前执行环境，包括执行模式等特性。
import mindspore.context as context
# c_transforms模块提供常用操作，包括OneHotOp和TypeCast
import mindspore.dataset.transforms as C
# vision.c_transforms模块是处理图像增强的高性能模块，用于数据增强图像数据改进训练模型。
import mindspore.dataset.vision as CV
import numpy as np
from mindspore import nn
from mindspore.nn import Accuracy
from mindspore.train import Model
from mindspore.train.callback import LossMonitor
import matplotlib.pyplot as plt
# 设置MindSpore的执行模式和设备
context.set_context(mode=context.GRAPH_MODE, device_target='CPU') # Ascend, CPU, GPU

对数据进行预处理：

def create_dataset(data_dir, training=True, batch_size=32, resize=(32, 32),
                   rescale=1/(255*0.3081), shift=-0.1307/0.3081, buffer_size=64):
    data_train = os.path.join(data_dir, 'train') # 训练集信息
    data_test = os.path.join(data_dir, 'test') # 测试集信息
    ds = ms.dataset.MnistDataset(data_train if training else data_test)
    ds = ds.map(input_columns=["image"], operations=[CV.Resize(resize), CV.Rescale(rescale, shift), CV.HWC2CHW()])
    ds = ds.map(input_columns=["label"], operations=C.TypeCast(ms.int32))
    ds = ds.shuffle(buffer_size=buffer_size).batch(batch_size, drop_remainder=True)
    return ds

MindSpore的model_zoo中提供了多种常见的模型，可以直接使用。我们构建LeNet-5模型：

class LeNet5(nn.Cell):
    def __init__(self):
        super(LeNet5, self).__init__()
        #设置卷积网络（输入输出通道数，卷积核尺寸，步长，填充方式）
        self.conv1 = nn.Conv2d(1, 6, 5, stride=1, pad_mode='valid')
        self.conv2 = nn.Conv2d(6, 16, 5, stride=1, pad_mode='valid')
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Dense(400, 120)
        self.fc2 = nn.Dense(120, 84)
        self.fc3 = nn.Dense(84, 10)
    #构建网络
    def construct(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool(x)
        x = self.relu(self.conv2(x))
        x = self.pool(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

使用MNIST数据集对上述定义的LeNet5模型进行训练。训练策略如下表所示，损失函数使用交叉熵损失函数。

batch size	number of epochs	learning rate	optimizer
32	5	0.01	Momentum 0.9

def train(data_dir, lr=0.01, momentum=0.9, num_epochs=5):
    ds_train = create_dataset(data_dir)
    ds_eval = create_dataset(data_dir, training=False)
    net = LeNet5()
    #计算softmax交叉熵。
    loss = nn.loss.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    #设置Momentum优化器
    opt = nn.Momentum(net.trainable_params(), lr, momentum)
    loss_cb = LossMonitor(per_print_times=ds_train.get_dataset_size())
    metrics = {"Accuracy": Accuracy(), "Confusion_matrix": nn.ConfusionMatrix(num_classes=10)}
    model = Model(net, loss, opt, metrics)
    model.train(num_epochs, ds_train, callbacks=[loss_cb], dataset_sink_mode=True)
    metrics_result = model.eval(ds_eval)
    res = metrics_result["Confusion_matrix"]
    print('Accuracy:',metrics_result["Accuracy"])
    print('Confusion_matrix:', res)
    return res

3.3.2 基于pytorch的LeNet-5网络

首先导入实验需要的模块：

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from matplotlib import pyplot as plt
from sklearn.metrics import confusion_matrix
from torchvision import datasets, transforms

接下来构建LeNet-5网络：

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5, 1 )
        self.conv2 = nn.Conv2d(6, 16, 5, 1)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(self.conv1(x), 2, 2)
        x = F.max_pool2d(self.conv2(x), 2, 2)
        x = x.view(-1, 16 * 4 * 4)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return F.log_softmax(x, dim=1)

训练和测试模块，采用训练策略与MindSpore保持一致，损失函数使用交叉熵损失函数。

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        pred = model(data)
        loss = F.nll_loss(pred, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if idx % 100 == 0:
            print("Train Epoch: {}, iterantion: {}, Loss: {}".format(epoch, idx, loss.item()))

def test(model, device, test_loader,):
    model.eval()
    total_loss = 0.
    correct = 0.
    predict=[]
    true=[]
    with torch.no_grad():
        for idx, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)
            output = model(data)
            total_loss += F.nll_loss(output, target, reduction="sum").item()
            pred = output.argmax(dim=1)
            true.append(target.tolist())
            predict.append(pred.tolist())
            correct += pred.eq(target.view_as(pred)).sum().item()

        total_loss /= len(test_loader.dataset)
        acc = correct / len(test_loader.dataset) * 100
        print("Test loss: {}, Accuracy: {}".format(total_loss, acc))
    predict = [n for a in predict for n in a]
    true=[n for a in true for n in a]
    return predict,true

4 实验结果及分析

4.1 实验结果

4.1.1 基于MindSpore的LeNet-5网络

4.1.1.1 准确率

可见经过5轮训练，准确率已达到98.267%。

4.1.1.2 混淆矩阵

混淆矩阵及其可视化效果如下：

4.1.2 基于pytorch的LeNet-5网络

4.1.2.1 准确率

可见经过5轮训练，准确率已超过97%。

4.1.2.2 混淆矩阵

4.2 结果分析与对比

用于手写数字识别（图像分类）的LeNet-5是卷积神经网络（CNN）的开山之作， AlexNet、ResNet等都是在其基础上发展而来的。CNN提出了三个创新点：局部感受野（Local Receptive Fields）、共享权重（Shared Weights）、时空下采样（Spatial or Temporal Subsampling）。

局部感受野（Local Receptive Fields）：用来表示网络内部的不同位置的神经元对原图像的感受范围的大小，对应于CNN中的卷积核，可以抽取图像初级的特征，如边、转角等，这些特征会在后来的层中通过各种联合的方式来检测高级别特征。

共享权重（Shared Weights）：在卷积过程中，每个卷积核所对应的窗口会以一定的步长在输入矩阵（图像）上不断滑动并进行卷积操作，最后，每个卷积核会生成一个对应的feature map（也就是卷积核的输出），一个feature map中的每一个单元都是由相同的权重（也就是对应的卷积核内的数值）计算得到的，这就是共享权重。

时空下采样（Spatial or Temporal Subsampling）：对应于CNN中的池化操作，也就是从卷积得到的feature map中提取出重要的部分，此操作可以降低模型对与图像平移和扭曲的敏感程度。

下表列出了各模型的准确率对比情况：

模型	准确率
基于MindSpore的LeNet-5网络	98.267%
基于pytorch的LeNet-5网络	97.35%

可见LeNet在不同平台的手写数字识别方面均展现出了优异的性能，无愧为卷积神经网络的开山鼻祖。

源代码

MSLeNet手写识别.py

import os
import mindspore as ms
# 导入mindspore中context模块，用于配置当前执行环境，包括执行模式等特性。
import mindspore.context as context
# c_transforms模块提供常用操作，包括OneHotOp和TypeCast
import mindspore.dataset.transforms as C
# vision.c_transforms模块是处理图像增强的高性能模块，用于数据增强图像数据改进训练模型。
import mindspore.dataset.vision as CV
import numpy as np
from mindspore import nn
from mindspore.nn import Accuracy
from mindspore.train import Model
from mindspore.train.callback import LossMonitor
import matplotlib.pyplot as plt
# 设置MindSpore的执行模式和设备
context.set_context(mode=context.GRAPH_MODE, device_target='CPU') # Ascend, CPU, GPU

def create_dataset(data_dir, training=True, batch_size=32, resize=(32, 32),
                   rescale=1/(255*0.3081), shift=-0.1307/0.3081, buffer_size=64):
    data_train = os.path.join(data_dir, 'train') # 训练集信息
    data_test = os.path.join(data_dir, 'test') # 测试集信息
    ds = ms.dataset.MnistDataset(data_train if training else data_test)
    ds = ds.map(input_columns=["image"], operations=[CV.Resize(resize), CV.Rescale(rescale, shift), CV.HWC2CHW()])
    ds = ds.map(input_columns=["label"], operations=C.TypeCast(ms.int32))
    ds = ds.shuffle(buffer_size=buffer_size).batch(batch_size, drop_remainder=True)
    return ds

ds = create_dataset('D:/Dataset/MNIST', training=False)
data = ds.create_dict_iterator().get_next()
images = data['image'].asnumpy()
labels = data['label'].asnumpy()
#显示前4张图片以及对应标签
for i in range(1, 5):
    plt.subplot(2, 2, i)
    plt.imshow(images[i][0])
    plt.title('Number: %s' % labels[i])
    plt.xticks([])
plt.show()

#定义LeNet5模型
class LeNet5(nn.Cell):
    def __init__(self):
        super(LeNet5, self).__init__()
        #设置卷积网络（输入输出通道数，卷积核尺寸，步长，填充方式）
        self.conv1 = nn.Conv2d(1, 6, 5, stride=1, pad_mode='valid')
        self.conv2 = nn.Conv2d(6, 16, 5, stride=1, pad_mode='valid')
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Dense(400, 120)
        self.fc2 = nn.Dense(120, 84)
        self.fc3 = nn.Dense(84, 10)
    #构建网络
    def construct(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool(x)
        x = self.relu(self.conv2(x))
        x = self.pool(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

def train(data_dir, lr=0.01, momentum=0.9, num_epochs=5):
    ds_train = create_dataset(data_dir)
    ds_eval = create_dataset(data_dir, training=False)
    net = LeNet5()
    #计算softmax交叉熵。
    loss = nn.loss.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    #设置Momentum优化器
    opt = nn.Momentum(net.trainable_params(), lr, momentum)
    loss_cb = LossMonitor(per_print_times=ds_train.get_dataset_size())
    metrics = {"Accuracy": Accuracy(), "Confusion_matrix": nn.ConfusionMatrix(num_classes=10)}
    model = Model(net, loss, opt, metrics)
    model.train(num_epochs, ds_train, callbacks=[loss_cb], dataset_sink_mode=True)
    metrics_result = model.eval(ds_eval)
    res = metrics_result["Confusion_matrix"]
    print('Accuracy:',metrics_result["Accuracy"])
    print('Confusion_matrix:', res)
    return res

# 绘制混淆矩阵
def plot_confusion_matrix(cm, title='Confusion Matrix'):
    classes = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
    plt.figure(figsize=(12, 8), dpi=100)
    np.set_printoptions(precision=2)
    # 混淆矩阵中每格的值
    ind_array = np.arange(len(classes))
    x, y = np.meshgrid(ind_array, ind_array)
    for x_val, y_val in zip(x.flatten(), y.flatten()):
        c = cm[y_val][x_val]
        if c > 0.001:
            plt.text(x_val, y_val, "%0.2f" % (c,), color='#EE3B3B', fontsize=10, va='center', ha='center')
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.binary)
    plt.title(title)
    plt.colorbar()
    xlocations = np.array(range(len(classes)))
    plt.xticks(xlocations, classes, rotation=90)
    plt.yticks(xlocations, classes)
    plt.ylabel('Actual Label')
    plt.xlabel('Predict Label')
    tick_marks = np.array(range(len(classes))) + 0.5
    plt.gca().set_xticks(tick_marks, minor=True)
    plt.gca().set_yticks(tick_marks, minor=True)
    plt.gca().xaxis.set_ticks_position('none')
    plt.gca().yaxis.set_ticks_position('none')
    plt.grid(True, which='minor', linestyle='-')
    plt.gcf().subplots_adjust(bottom=0.15)
    plt.show()

confusion = train('D:/Dataset/MNIST')
plot_confusion_matrix(confusion,title='Confusion Matrix')

pytorchLeNet手写识别.py

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from matplotlib import pyplot as plt
from sklearn.metrics import confusion_matrix
from torchvision import datasets, transforms

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5, 1 )
        self.conv2 = nn.Conv2d(6, 16, 5, 1)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(self.conv1(x), 2, 2)
        x = F.max_pool2d(self.conv2(x), 2, 2)
        x = x.view(-1, 16 * 4 * 4)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return F.log_softmax(x, dim=1)


def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        pred = model(data)
        loss = F.nll_loss(pred, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # if idx % 100 == 0:
        #     print("Train Epoch: {}, iterantion: {}, Loss: {}".format(epoch, idx, loss.item()))

def test(model, device, test_loader,):
    model.eval()
    total_loss = 0.
    correct = 0.
    predict=[]
    true=[]
    with torch.no_grad():
        for idx, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)
            output = model(data)
            total_loss += F.nll_loss(output, target, reduction="sum").item()
            pred = output.argmax(dim=1)
            true.append(target.tolist())
            predict.append(pred.tolist())
            correct += pred.eq(target.view_as(pred)).sum().item()

        total_loss /= len(test_loader.dataset)
        acc = correct / len(test_loader.dataset) * 100
        print("Test loss: {}, Accuracy: {}".format(total_loss, acc))
    predict = [n for a in predict for n in a]
    true=[n for a in true for n in a]
    return predict,true

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size = 32
train_dataloader = torch.utils.data.DataLoader(
    datasets.MNIST("D:/Dataset/pytorch/", train=True, download=False,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))

                   ])),
    batch_size=batch_size, shuffle=True,
    num_workers=1, pin_memory=True  # True加快训练
)
test_dataloader = torch.utils.data.DataLoader(
    datasets.MNIST("D:/Dataset/pytorch/", train=False, download=False,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=batch_size, shuffle=True,
    num_workers=1, pin_memory=True
)

# 绘制混淆矩阵
def plot_confusion_matrix(cm, title='Confusion Matrix'):
    classes = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
    plt.figure(figsize=(12, 8), dpi=100)
    np.set_printoptions(precision=2)
    # 混淆矩阵中每格的值
    ind_array = np.arange(len(classes))
    x, y = np.meshgrid(ind_array, ind_array)
    for x_val, y_val in zip(x.flatten(), y.flatten()):
        c = cm[y_val][x_val]
        if c > 0.001:
            plt.text(x_val, y_val, "%0.2f" % (c,), color='#EE3B3B', fontsize=10, va='center', ha='center')
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.binary)
    plt.title(title)
    plt.colorbar()
    xlocations = np.array(range(len(classes)))
    plt.xticks(xlocations, classes, rotation=90)
    plt.yticks(xlocations, classes)
    plt.ylabel('Actual Label')
    plt.xlabel('Predict Label')
    tick_marks = np.array(range(len(classes))) + 0.5
    plt.gca().set_xticks(tick_marks, minor=True)
    plt.gca().set_yticks(tick_marks, minor=True)
    plt.gca().xaxis.set_ticks_position('none')
    plt.gca().yaxis.set_ticks_position('none')
    plt.grid(True, which='minor', linestyle='-')
    plt.gcf().subplots_adjust(bottom=0.15)
    plt.show()

if __name__ == '__main__':
    lr = 0.01
    momentum = 0.9
    model = LeNet().to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
    num_eopchs = 5
    for eopch in range(num_eopchs):
        train(model, device, train_dataloader, optimizer, eopch)
        y_pre,y=test(model, device, test_dataloader)
        plot_confusion_matrix(confusion_matrix(y, y_pre), title='Confusion Matrix')