详解PyTorch手写数字识别(MNIST数据集)

MNIST 手写数字识别是一个比较简单的入门项目,相当于深度学习中的 Hello World,可以让我们快速了解构建神经网络的大致过程。虽然网上的案例比较多,但还是要自己实现一遍。代码采用 PyTorch 1.0 编写并运行。

导入相关库

import torch

import torch.nn as nn

import torch.nn.functional as F

import torch.optim as optim

from torchvision import datasets, transforms

import torchvision

from torch.autograd import Variable

from torch.utils.data import DataLoader

import cv2

torchvision 用于下载并导入数据集

cv2 用于展示数据的图像

获取训练集和测试集

# 下载训练集

train_dataset = datasets.MNIST(root='./num/',

train=True,

transform=transforms.ToTensor(),

download=True)

# 下载测试集

test_dataset = datasets.MNIST(root='./num/',

train=False,

transform=transforms.ToTensor(),

download=True)

root 用于指定数据集在下载之后的存放路径

transform 用于指定导入数据集需要对数据进行那种变化操作

train是指定在数据集下载完成后需要载入的那部分数据,设置为 True 则说明载入的是该数据集的训练集部分,设置为 False 则说明载入的是该数据集的测试集部分

download 为 True 表示数据集需要程序自动帮你下载

这样设置并运行后,就会在指定路径中下载 MNIST 数据集,之后就可以使用了。

数据装载和预览

# dataset 参数用于指定我们载入的数据集名称

# batch_size参数设置了每个包中的图片数据个数

# 在装载的过程会将数据随机打乱顺序并进打包

# 装载训练集

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,

batch_size=batch_size,

shuffle=True)

# 装载测试集

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,

batch_size=batch_size,

shuffle=True)

在装载完成后,可以选取其中一个批次的数据进行预览:

images, labels = next(iter(data_loader_train))

img = torchvision.utils.make_grid(images)

img = img.numpy().transpose(1, 2, 0)

std = [0.5, 0.5, 0.5]

mean = [0.5, 0.5, 0.5]

img = img * std + mean

print(labels)

cv2.imshow('win', img)

key_pressed = cv2.waitKey(0)

在以上代码中使用了 iter 和 next 来获取取一个批次的图片数据和其对应的图片标签,然后使用 torchvision.utils 中的 make_grid 类方法将一个批次的图片构造成网格模式。

预览图片如下:


并且打印出了图片相对应的数字:

搭建神经网络

# 卷积层使用 torch.nn.Conv2d

# 激活层使用 torch.nn.ReLU

# 池化层使用 torch.nn.MaxPool2d

# 全连接层使用 torch.nn.Linear

class LeNet(nn.Module):

def __init__(self):

super(LeNet, self).__init__()

self.conv1 = nn.Sequential(nn.Conv2d(1, 6, 3, 1, 2), nn.ReLU(),

nn.MaxPool2d(2, 2))

self.conv2 = nn.Sequential(nn.Conv2d(6, 16, 5), nn.ReLU(),

nn.MaxPool2d(2, 2))

self.fc1 = nn.Sequential(nn.Linear(16 * 5 * 5, 120),

nn.BatchNorm1d(120), nn.ReLU())

self.fc2 = nn.Sequential(

nn.Linear(120, 84),

nn.BatchNorm1d(84),

nn.ReLU(),

nn.Linear(84, 10))

# 最后的结果一定要变为 10,因为数字的选项是 0 ~ 9

def forward(self, x):

x = self.conv1(x)

x = self.conv2(x)

x = x.view(x.size()[0], -1)

x = self.fc1(x)

x = self.fc2(x)

x = self.fc3(x)

return x

前向传播内容:

首先经过 self.conv1() 和 self.conv1() 进行卷积处理

然后进行 x = x.view(x.size()[0], -1),对参数实现扁平化(便于后面全连接层输入)

最后通过 self.fc1() 和 self.fc2() 定义的全连接层进行最后的分类

训练模型

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

batch_size = 64

LR = 0.001

net = LeNet().to(device)

# 损失函数使用交叉熵

criterion = nn.CrossEntropyLoss()

# 优化函数使用 Adam 自适应优化算法

optimizer = optim.Adam(

net.parameters(),

lr=LR,

)

epoch = 1

if __name__ == '__main__':

for epoch in range(epoch):

sum_loss = 0.0

for i, data in enumerate(train_loader):

inputs, labels = data

inputs, labels = Variable(inputs).cuda(), Variable(labels).cuda()

optimizer.zero_grad() #将梯度归零

outputs = net(inputs) #将数据传入网络进行前向运算

loss = criterion(outputs, labels) #得到损失函数

loss.backward() #反向传播

optimizer.step() #通过梯度做一步参数更新

# print(loss)

sum_loss += loss.item()

if i % 100 == 99:

print('[%d,%d] loss:%.03f' %

(epoch + 1, i + 1, sum_loss / 100))

sum_loss = 0.0

测试模型

net.eval() #将模型变换为测试模式

correct = 0

total = 0

for data_test in test_loader:

images, labels = data_test

images, labels = Variable(images).cuda(), Variable(labels).cuda()

output_test = net(images)

_, predicted = torch.max(output_test, 1)

total += labels.size(0)

correct += (predicted == labels).sum()

print("correct1: ", correct)

print("Test acc: {0}".format(correct.item() /

len(test_dataset)))

训练及测试的情况:


98% 以上的成功率,效果还不错。

以上就是本文的全部内容,希望对大家的学习有所帮助,也希望大家多多支持。

以上是 详解PyTorch手写数字识别(MNIST数据集) 的全部内容, 来源链接: utcz.com/z/350791.html

回到顶部