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

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

---

---

前言

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

---

看本文之前，需掌握基础的深度学习知识（包括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
1
2
3
4
5
6
7
8

2.设置相关参数

epochs = 10
batch_size = 64
lr = 0.001
1
2
3

这里详细解释一下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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

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

对于transform：

对于datasets.MNIST：

对于DataLoader：

预览（可选）

# 数据预览
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()
1
2
3
4
5
6
7
8
9
10

效果：

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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

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

5.训练

相关解释见代码注释。

# 训练和参数优化

# 定义求导函数
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)))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

6.保存模型

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

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

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')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

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

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

---

三、其他

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

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