【Pytorch】MNIST 图像分类代码 - 超详细解读

最近机器学习在低年级本科生中热度剧增，小编经常看见在自习室里啃书的小伙伴。但是缺少经验指导，也许原理清楚了，但是由于很多书中对细节上的函数等等介绍不多，很多小伙伴对于具体代码只是一知半解。这篇文章基于当下最热门的学习框架 Pytorch，详细讲解图像分类中最基础的图像分类 —— MNIST 数据集分类。

千鱼干

12199人浏览 · 2021-12-07 00:41:48

千鱼干 · 2021-12-07 00:41:48 发布

【Pytorch】MNIST 图像分类代码 - 超详细解读

前言

最近机器学习在低年级本科生中热度剧增，小编经常看见在自习室里啃相关书籍的小伙伴。但由于缺少经验指导，也许原理清楚了，但是由于很多书中对细节上的函数等等介绍不多，很多小伙伴对于具体代码只是一知半解。这篇文章基于当下最热门的学习框架 Pytorch，详细讲解图像分类中最基础的图像分类 —— MNIST 数据集分类。
同时，希望这篇文章会帮助大家领会基本的深度学习思路。

看本文之前，需掌握基础的深度学习知识（包括CNN（卷积神经网络））。如没有基础或者基础薄弱，请至少先按照此表格顺序熟悉一下相关内容：

推荐顺序	网址
1	如何自己从零实现一个神经网络? 量子位的回答 - 知乎
2	卷积神经网络CNN完全指南终极版（一）沉迷学习的糕糕的文章 - 知乎
3	卷积神经网络CNN完全指南终极版（二）沉迷学习的糕糕的文章 - 知乎

熟悉以上文章之后，基本就可以开始阅读这篇博客啦！

文章中遇到任何关于Pytorch框架代码的问题，可先查询快速手册 - Pytorch官网相关内容。如未找到，请在 Pytorch官网查询。

一、代码框架

下面是我本人比较喜欢的代码框架，可以参考。

文件名：model.py

1.引入包
2.设置相关参数
3.处理数据集
——定义transform
——导入数据集
——装载（DataLoader）
——预览（可选）
4.构建网络
5.训练
6.保存模型

二、实现代码

1.引入包

代码如下：

import torch
import torch.nn as nn
from torch.nn import Sequential
from matplotlib import pyplot as plt
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
from torchvision.utils import make_grid
import torchvision.transforms as transforms

包名	功能
torch	核心包
torch.nn	包含神经网络的Modules和用来继承的包以及一些函数方法（nn.functional）
torchvision	包含一些数据集、模型、图像处理方法
torch.utils	一个工具包
matplotlib	用于显示数据集图片

2.设置相关参数

epochs = 10
batch_size = 64
lr = 0.001

参数	意义
epochs	被训练几轮
batch_size	每批次大小，即每轮迭代训练时每次的数据量
lr	learning rate，即学习率。一般用很小的值

这里详细解释一下epochs和batch_szie：
->batch_size表示每轮迭代训练时每次训练的数据量；
->epochs表示训练几轮。

每一次迭代（Iteration）都是一次权重更新，每一次权重更新需要batch_size个数据进行正向传递（Forward）运算得到损失函数，再通过反向传导（Backward）更新参数（注意，在这个过程中需要把梯度（Grad）设置为0，这个后面再讲）。1个迭代等于使用batch_size个样本训练一次。比如有256个样本数据，完整训练完这些样本数据需要：
->batch_size=64；
->迭代4次；
->epochs=1。

而通常会将epochs设为不仅1次，这就跟磨面一样，磨完一轮不够，磨多轮才能得到更加精细的面粉。

3.处理数据集

# 设置数据转换方式
transform = transforms.Compose([
    transforms.ToTensor(),  # 把数据转换为张量（Tensor）
    transforms.Normalize(  # 标准化，即使数据服从期望值为 0，标准差为 1 的正态分布
        mean=[0.5, ],  # 期望
        std=[0.5, ]  # 标准差
    )
])

# 训练集导入
data_train = datasets.MNIST(root='data/', transform=transform, train=True, download=True)
# 数据集导入
data_test = datasets.MNIST(root='data/', transform=transform, train=False)

# 数据装载
# 训练集装载
dataloader_train = DataLoader(dataset=data_train, batch_size=64, shuffle=True)
# 数据集装载
dataloader_test = DataLoader(dataset=data_test, batch_size=64, shuffle=True)

除了 代码内的注释 之外，在这段代码中一些方法或参数的解释如下。

对于transform：

参数	意义
transforms.ToTensor()	把数据转换为张量（Tensor）
transforms.Normalize	标准化，即使数据服从期望值为 0，标准差为 1 的正态分布
mean	期望
std	标准差

对于datasets.MNIST：

参数	意义
root	数据集（此处为MNIST）路径
transform	转换形式
train	是否训练。对于训练集，train=True，对于测试集，train=False
download	是否下载（会自动判断是否下载过或数据集是否存在于root下，是的话再次训练时就不下载了）

对于DataLoader：

参数	意义
dataset	要处理的数据集
batch_size	批次大小
shuffle	是否打乱数据顺序

预览（可选）

# 数据预览
images, labels = next(iter(dataloader_train))
img = make_grid(images)
img = img.numpy().transpose(1, 2, 0)
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]
img = img * std + mean
print([labels[i] for i in range(16)])
plt.imshow(img)
plt.show()

方法	作用
iter(dataloader_train)	生成dataloader_train的迭代器
next	返回迭代器的下一个项目（配合iter()使用）
make_grid	生成网格
img.numpy().transpose(1, 2, 0)	将img的numpy数组矩阵的C、W、H位置调换。括号内的1, 2, 0表示将原来第1, 2, 0位置换位0, 1, 2（即把原本[C, W, H]矩阵转换为[H, W, C]矩阵。Pytorch中使用的数据格式与plt.imshow()函数的格式不一致，Pytorch中为[C, H, W]，而plt.imshow()中则是[H, W, C]。其中C=Channel，即颜色通道；H=Height，图像长度；Width，图片宽度）
plt.imshow(img)和plt.show()	显示图片

效果：

4.构建网络

# 构建卷积神经网络
class CNN(nn.Module):  # 从父类 nn.Module 继承
    def __init__(self):  # 相当于 C++ 的构造函数
        # super() 函数是用于调用父类(超类)的一个方法，是用来解决多重继承问题的
        super(CNN, self).__init__()
        
        # 第一层卷积层。Sequential(意为序列) 括号内表示要进行的操作
        self.conv1 = Sequential(
            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        # 第二卷积层
        self.conv2 = Sequential(
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        # 全连接层（Dense，密集连接层）
        self.dense = Sequential(
            nn.Linear(7 * 7 * 128, 1024),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(1024, 10)
        )

    def forward(self, x):  # 正向传播
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        x = x2.view(-1, 7 * 7 * 128)
        x = self.dense(x)
        return x

除了 代码内的注释 之外，在这段代码中一些方法或参数的解释如下：

方法或参数	意义或作用
nn.Conv2d	对二维图像的卷积操作。其中`in_channels`代表输入通道，`out_channels`代表输出通道，`kernel_size`代表卷积核大小（n * n），`stride`代表卷积核移动的步长，`padding`代表填充大小（属于基础内容。具体内容请自行百度）
nn.BatchNorm2d	Batch Normalization（BN），批标准化。使一批次特征图（Feature Map）满足均值为0，方差为1的正态分布。作用：加速收敛；控制过拟合，以少用或不用Dropout和正则；降低网络对初始化权重不敏感；允许使用较大的学习率
nn.ReLU	一种常用激活函数，不作赘述
nn.MaxPool2d	对二维图像做最大池化处理，不做赘述
nn.Linear	不再赘述
nn.Dropout	Dropout，防止过拟合，不做赘述
x2.view(-1, 7 * 7 * 128)	参数扁平化，使全连接层输出的参数维度和其输入维度匹配

5.训练

6.保存模型

torch.save(cnn, 'data/model.pth')  # 将模型保存到当前目录下 data 文件夹内，名为 model.pth

恭喜！如果你做到了这一步，训练的所有步骤就完成啦！

MNIST 图像识别完整代码如下：

import torch
import torch.nn as nn
from torch.nn import Sequential
from matplotlib import pyplot as plt
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
from torchvision.utils import make_grid
import torchvision.transforms as transforms

epochs = 10
batch_size = 64
lr = 0.001

# 转换导集

# 设置数据转换方式
transform = transforms.Compose([
    transforms.ToTensor(),  # 把数据转换为张量（Tensor）
    transforms.Normalize(  # 标准化，即使数据服从期望值为 0，标准差为 1 的正态分布
        mean=[0.5, ],  # 期望
        std=[0.5, ]  # 标准差
    )
])

# 训练集导入
data_train = datasets.MNIST(root='data/', transform=transform, train=True, download=True)
# 数据集导入
data_test = datasets.MNIST(root='data/', transform=transform, train=False)

# 数据装载

# 训练集装载
dataloader_train = DataLoader(dataset=data_train, batch_size=64, shuffle=True)
# 数据集装载
dataloader_test = DataLoader(dataset=data_test, batch_size=64, shuffle=True)

# 数据预览
images, labels = next(iter(dataloader_train))
img = make_grid(images)
img = img.numpy().transpose(1, 2, 0)
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]
img = img * std + mean
print([labels[i] for i in range(16)])
plt.imshow(img)
plt.show()


# 构建卷积神经网络
class CNN(nn.Module):  # 从父类 nn.Module 继承
    def __init__(self):  # 相当于 C++ 的构造函数
        # super() 函数是用于调用父类(超类)的一个方法，是用来解决多重继承问题的
        super(CNN, self).__init__()

        # 第一层卷积层。Sequential(意为序列) 括号内表示要进行的操作
        self.conv1 = Sequential(
            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        # 第二卷积层
        self.conv2 = Sequential(
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        # 全连接层（Dense，密集连接层）
        self.dense = Sequential(
            nn.Linear(7 * 7 * 128, 1024),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(1024, 10)
        )

    def forward(self, x):  # 正向传播
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        x = x2.view(-1, 7 * 7 * 128)
        x = self.dense(x)
        return x


# 训练和参数优化

# 定义求导函数
def get_Variable(x):
    x = torch.autograd.Variable(x)  # Pytorch 的自动求导

    # 判断是否有可用的 GPU
    return x.cuda() if torch.cuda.is_available() else x


# 定义网络
cnn = CNN()

# 判断是否有可用的 GPU 以加速训练
if torch.cuda.is_available():
    cnn = cnn.cuda()

# 设置损失函数为 CrossEntropyLoss（交叉熵损失函数）
loss_F = nn.CrossEntropyLoss()

# 设置优化器为 Adam 优化器
optimizer = torch.optim.Adam(cnn.parameters(), lr=lr)

# 训练
for epoch in range(epochs):
    running_loss = 0.0  # 一个 epoch 的损失
    running_correct = 0.0  # 准确率
    print("Epoch [{}/{}]".format(epoch, epochs))
    for data in dataloader_train:
        # DataLoader 返回值是一个 batch 内的图像和对应的 label
        X_train, y_train = data
        X_train, y_train = get_Variable(X_train), get_Variable(y_train)
        outputs = cnn(X_train)
        _, pred = torch.max(outputs.data, 1)
        # 后面的参数代表降低 outputs.data 的维度 1 个维度再输出
        # 第一个返回值是张量中最大值，第二个返回值是最大值索引
        # -------------------下面内容与随机梯度下降类似-----------------------------
        optimizer.zero_grad()
        # 梯度置零
        loss = loss_F(outputs, y_train)
        # 求损失
        loss.backward()
        # 反向传播
        optimizer.step()
        # 更新所有梯度
        # --------------------上面内容与随机梯度下降类似----------------------------
        running_loss += loss.item()  # 此处 item() 表示返回每次的 loss 值
        running_correct += torch.sum(pred == y_train.data)

    testing_correct = 0.0

    for data in dataloader_test:
        X_test, y_test = data
        X_test, y_test = get_Variable(X_test), get_Variable(y_test)
        outputs = cnn(X_test)
        _, pred = torch.max(outputs, 1)
        testing_correct += torch.sum(pred == y_test.data)
        # print(testing_correct)
    print("Loss: {:.4f}  Train Accuracy: {:.4f}%  Test Accuracy: {:.4f}%".format(
        running_loss / len(data_train), 100 * running_correct / len(data_train),
        100 * testing_correct / len(data_test)))

# 保存模型
torch.save(cnn, 'data/model.pth')

注：在inference.py内加载模型时：

# 加载模型
cnn = torch.load('data/model.pth')
cnn.eval()  # 进入推断模式

三、其他

作者是某高校大二学生，计算机科学与技术在读。大一下学期接触机器学习，之前主攻超分辨率重构。机器学习纯属业余爱好，几乎无人指导，故文章若有纰漏，望批评指正！

*本博客部分内容来源于网络。

华为开发者空间

华为开发者空间，是为全球开发者打造的专属开发空间，汇聚了华为优质开发资源及工具，致力于让每一位开发者拥有一台云主机，基于华为根生态开发、创新。

更多推荐

GaussDB Ustore存储引擎解读

华为开发者空间

如何在鲲鹏平台上快速上手应用开发？鲲鹏DevKit给你答案

鲲鹏DevKit针对不同的业务场景，提供了应用迁移和系统迁移两套解决方案，帮忙开发者快速从X86平台迁移至鲲鹏平台，通过详细的迁移建议降低迁移门槛，可视化展示迁移进度，打消鲲鹏平台开发的顾虑。

华为开发者空间

华为云开源项目Sermant正式成为CNCF官方项目

华为开发者空间

所有评论(0)

查看更多评论

千鱼干

@qq_25426559

已为社区贡献1条内容