1. Johngo学长首页
  2. 技术点详细复盘专题
  3. 人工智能

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

• 人工智能 • 阅读 77

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

目录

前言

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

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

推荐顺序网址1
如何自己从零实现一个神经网络? 量子位的回答 – 知乎 卷积神经网络CNN完全指南终极版（一） 沉迷学习的糕糕的文章 – 知乎 卷积神经网络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每批次大小，即每轮迭代训练时每次的数据量lrlearning 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(),
    transforms.Normalize(
        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=Falsedownload是否下载（会自动判断是否下载过或数据集是否存在于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()显示图片

效果：

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

4.构建网络


class CNN(nn.Module):
    def __init__(self):

        super(CNN, self).__init__()

        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)
        )

        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.BatchNorm2dBatch Normalization（BN），批标准化。使一批次特征图（Feature Map）满足均值为0，方差为1的正态分布。作用：加速收敛；控制过拟合，以少用或不用Dropout和正则；降低网络对初始化权重不敏感；允许使用较大的学习率nn.ReLU一种常用激活函数，不作赘述nn.MaxPool2d对二维图像做最大池化处理，不做赘述nn.Linear不再赘述nn.DropoutDropout，防止过拟合，不做赘述x2.view(-1, 7 * 7 * 128)参数扁平化，使全连接层输出的参数维度和其输入维度匹配

5.训练

相关解释见代码注释。


def get_Variable(x):
    x = torch.autograd.Variable(x)

    return x.cuda() if torch.cuda.is_available() else x

cnn = CNN()

if torch.cuda.is_available():
    cnn = cnn.cuda()

loss_F = nn.CrossEntropyLoss()

optimizer = torch.optim.Adam(cnn.parameters(), lr=lr)

for epoch in range(epochs):
    running_loss = 0.0
    running_correct = 0.0
    print("Epoch [{}/{}]".format(epoch, epochs))
    for data in dataloader_train:

        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)

        optimizer.zero_grad()

        loss = loss_F(outputs, y_train)

        loss.backward()

        optimizer.step()

        running_loss += loss.item()
        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("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)))

6.保存模型

torch.save(cnn, '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(),
    transforms.Normalize(
        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):
    def __init__(self):

        super(CNN, self).__init__()

        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)
        )

        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)

    return x.cuda() if torch.cuda.is_available() else x

cnn = CNN()

if torch.cuda.is_available():
    cnn = cnn.cuda()

loss_F = nn.CrossEntropyLoss()

optimizer = torch.optim.Adam(cnn.parameters(), lr=lr)

for epoch in range(epochs):
    running_loss = 0.0
    running_correct = 0.0
    print("Epoch [{}/{}]".format(epoch, epochs))
    for data in dataloader_train:

        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)

        optimizer.zero_grad()

        loss = loss_F(outputs, y_train)

        loss.backward()

        optimizer.step()

        running_loss += loss.item()
        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("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()

三、其他

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

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

Original: https://blog.csdn.net/qq_25426559/article/details/121712992
Author: 千鱼干
Title: 【Pytorch】MNIST 图像分类代码 – 超详细解读

原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/705939/

转载文章受原作者版权保护。转载请注明原作者出处！

赞 (0)
0

【自取】最近整理的，有需要可以领取学习：



Linux核心资料大放送~

全栈面试题汇总（持续更新&可下载）

一个提高学习100%效率的工具！

【超详细】深度学习面试题目！

LeetCode Python刷题答案下载！

LeetCode Java版刷题答案下载！

LeetCode C++ 版本，抓紧保存！

LeetCode GO语言 刷题答案下载！

大家都在看

亲爱的 Coder【最近整理，可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球