实现一个手写数字识别项目

发表于 更新于
MNIST手写数字识别项目因为数据量小识别任务简单而成为图像识别入门的第一课。

MNIST手写数字识别项目因为数据量小识别任务简单而成为图像识别入门的第一课。MNIST数据集包括6万张28x28的训练样本，1万张测试样本。

1、数据查看

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
import numpy as np
import matplotlib.pyplot as plt
import torch
import torchvision as tv
from torchvision import transforms

#========================================================
#  超参数设置
#========================================================

# 批处理尺寸(batch_size)
BATCH_SIZE = 20   
# 定义图像类别
classes = ['0', '1', '2', '3', '4', 
           '5', '6', '7', '8', '9']

#========================================================
#  数据加工
#========================================================

# 定义数据预处理方式
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图片(Image)转化为Tensor, 归一化到[0,1]
    transforms.Normalize(mean=[0.5], std=[0.5])  # 标准化到[-1,1]
    ])

# 定义训练数据集
trainset = tv.datasets.MNIST(
    root='./data/',
    train=True,
    download=True,
    transform=transform)

# 定义训练批处理数据
trainloader = torch.utils.data.DataLoader(
    trainset,
    batch_size=BATCH_SIZE,
    shuffle=True
    )

#========================================================
#  查看原始图像
#========================================================
    
# 获取一个批次的训练数据
dataiter = iter(trainloader)
images, labels = dataiter.next()
images = images.numpy()

# 打印图片配上相应的标签
fig = plt.figure(figsize=(25, 4))
for idx in np.arange(BATCH_SIZE):
    ax = fig.add_subplot(2, BATCH_SIZE/2, idx+1, xticks=[], yticks=[])
    ax.imshow(np.squeeze(images[idx]), cmap='gray')
    ax.set_title(classes[labels[idx]])
plt.show()

#========================================================
#  查看归一化后的图像
#========================================================

idx = 2
img = np.squeeze(images[idx])

fig = plt.figure(figsize = (12,12)) 
ax = fig.add_subplot(111)
ax.imshow(img, cmap='gray')
width, height = img.shape
thresh = img.max()/2.5
for x in range(width):
    for y in range(height):
        val = round(img[x][y],2) if img[x][y] !=0 else 0
        ax.annotate(str(val), xy=(y,x),
                    horizontalalignment='center',
                    verticalalignment='center',
                    color='white' if img[x][y]<thresh else 'black')
plt.show()

2、完整代码

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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
import numpy as np
import torch
from torch import nn
import torch.optim as optim
import torchvision as tv
import torchvision.transforms as transforms
from matplotlib import pyplot as plt
import random

#========================================================
#  超参数设置和类别配置
#========================================================

EPOCH = 5   # 遍历数据集次数
BATCH_SIZE = 64   # 批处理尺寸(batch_size)
LR = 0.001    # 学习率
# 定义图像类别
classes = ['0', '1', '2', '3', '4', 
           '5', '6', '7', '8', '9']

#========================================================
#  数据加工
#========================================================

# 定义数据预处理方式
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图片(Image)转化为Tensor, 归一化到[0,1]
    transforms.Normalize(mean=[0.5], std=[0.5])  # 标准化到[-1,1]
    ])

# 定义训练数据集
trainset = tv.datasets.MNIST(
    root='./data/',
    train=True,
    download=True,
    transform=transform)

# 定义训练批处理数据
trainloader = torch.utils.data.DataLoader(
    trainset,
    batch_size=BATCH_SIZE,
    shuffle=True
    )

# 定义测试数据集
testset = tv.datasets.MNIST(
    root='./data/',
    train=False,
    download=True,
    transform=transform)

# 定义测试批处理数据
testloader = torch.utils.data.DataLoader(
    testset,
    batch_size=BATCH_SIZE,
    shuffle=False
    )

#========================================================
#  模型结构设计
#========================================================

class LeNet5(nn.Module):
    '''
    LeNet5网络
    INPUT  -> 图像规格(28, 28, 1), 待分类数(10)
    '''
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 6, kernel_size=5, padding=2, bias=False),
            # 第一次卷积后图像尺寸 = (28+2*2-5)/步长+1 = 28
            nn.Tanh(),
            nn.MaxPool2d(kernel_size=2, stride=2)
            # 经过池化层后图像尺寸 = (28-2)/步长+1 = 14
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6, 16, kernel_size=5, bias=False),
            # 第二次卷积后图像尺寸 = (14-5)/步长+1 = 10
            nn.Tanh(),
            nn.MaxPool2d(kernel_size=2, stride=2)
            # 经过池化层后图像尺寸 = (10-2)/步长+1 = 5
        )
        self.classifier = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.Tanh(),
            nn.Linear(120, 84),
            nn.Tanh(),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        # nn.Linear()的输入输出都是一维数组，所以要把tensor展成一维
        x = x.view(x.size(0), 16*5*5)
        # x = x.view(x.size()[0], -1)
        x = self.classifier(x)
        return x

#========================================================
#  模型实例化
#========================================================

# 定义是否使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = LeNet5().to(device)

# 定义损失函数loss function 和优化方式（采用SGD）
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数，通常用于多分类问题上
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9)

#========================================================
#  训练模型和测试模型
#========================================================

def train_work():
    '''
    训练阶段
    '''
    net.train()  # 设置模式为训练模式
    loss_over_time = []
    for epoch in range(EPOCH):
        running_loss = 0.0
        # 数据读取
        i = 0
        for data in trainloader:
            i += 1
            images, labels = data
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            
            # forward + backward + optimize
            outputs = net(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            # 每训练100个batch打印一次平均loss
            running_loss += loss.item()
            if i % 100 == 99:
                avg_loss = running_loss/100
                loss_over_time.append(avg_loss)
                print('[%d] loss: %.03f' % (epoch + 1, avg_loss))
                running_loss = 0.0

        # 每跑完一次epoch测试一下准确率
        with torch.no_grad():
            correct = 0
            total = 0
            for data in testloader:
                images, labels = data
                images, labels = images.to(device), labels.to(device)
                outputs = net(images)
                # 取得分最高的那个类
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum()
            print('第%d个epoch的预测准确率为：%d%%' % (epoch + 1, (100 * correct / total)))
    torch.save(net.state_dict(), '%s/net_best.pth' % ('./model/'))
    print('Finished Training')
    return loss_over_time

def test_work():
    '''
    测试阶段
    '''
    net.load_state_dict(torch.load("./model/net_best.pth")) # 加载模型参数
    net.eval()  # 设置模式为验证模式
    class_correct = list(0. for i in range(len(classes)))
    class_total = list(0. for i in range(len(classes)))
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = net(images)
        # 取得分最高的那个类
        _, predicted = torch.max(outputs.data, 1)
        correct = np.squeeze(predicted.eq(labels.data.view_as(predicted)))
        # 统计每一类的正确个数
        for i in range(len(data)):
            label = labels.data[i]
            class_correct[label] += correct[i].item()
            class_total[label] += 1
    # 计算每一类的准确率
    for i in range(len(classes)):
        if class_total[i] > 0:
            print('%5s的识别准确率 : %2d%% (%2d/%2d)' % (classes[i], 100 * class_correct[i] / class_total[i],
                                                        np.sum(class_correct[i]), np.sum(class_total[i])))
    print('当前模型预测准确率为：%d%% (%2d/%2d)' % (100 * np.sum(class_correct) / np.sum(class_total),
                                                np.sum(class_correct), np.sum(class_total)))

def show_loss(training_loss):
    '''
    可视化损失变化
    '''
    plt.plot(training_loss)
    plt.xlabel('100\'s of batches')
    plt.ylabel('loss')
    plt.ylim(0, 2.5) # consistent scale
    plt.show()

def show_predicted():
    '''
    展示预测结果
    '''
    net.load_state_dict(torch.load("./model/net_best.pth"))
    net.eval()  # 设置模式为验证模式
    n_examples = 12
    [examples] = random.sample(list(testloader), 1)
    images, labels = examples
    images, labels = images.to(device), labels.to(device)
    outputs = net(images)
    _, predicted = torch.max(outputs.data, 1)
    
    for i in range(n_examples):
        img = images[i,0,:,:]
        img = img.cpu()
        img = img.data.numpy()

        img = 1.0/(1+np.exp(-1*img))
        img = np.round(img*255)
        
        ax = plt.subplot((n_examples//6), 6, i+1)
        plt.imshow(img, cmap='gray')
        plt.title('label:{}\nPredicted: {}'.format(labels[i],predicted[i]),
                 color=("green" if labels[i]==predicted[i] else "red"))
        plt.axis('off')
    plt.show()

#========================================================
#  主程序
#========================================================

if __name__ == "__main__":
    training_loss = train_work()
    show_loss(training_loss)
    test_work()
    show_predicted()

读入单张图片验证

1