最终成果
GITHUB
本文以最经典的mnist
数据集为例,讲述了使用pytorch
做机器学习的一整套流程,文中所提到的所有代码都可以到github
中查看。
项目场景
简单的学习pytorch
、自动求导和神经网络的知识后,我们来练习使用mnist
数据集训练一个cnn
手写数字识别模型。
导入模块
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
pytorch
安装教程可以看这这两篇文章:Windows 和 Linux
matplotlib
库用于绘图,如果没有请通过pip
命令安装:
pip install matplotlib
训练设备
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
如果有可用的gpu
就在gpu
中训练,没有就使用cpu
训练。
定义超参数
EPOCH = 10
BATCH_SIZE = 128
LR = 1E-3
EPOCH
:训练的轮数。
BATCH_SIZE
:数据加载器的批次大小。
LR
:优化器的学习率。
下载数据集
train_file = datasets.MNIST(
root='./dataset/',
train=True,
transform=transforms.ToTensor(),
download=True
)
test_file = datasets.MNIST(
root='./dataset/',
train=False,
transform=transforms.ToTensor()
)
root
:指定数据集的存放路径。
train
:True
表示训练集;False
表示测试集。
transform
:采用何种方式进行图像转换;transforms.ToTensor()
是将形状为(H x W x C)
的图片数据转换为形状为(C x H x W)
的张量,C
为图片通道数,这里是1,也就是灰度图片,H
、W
分别为图片宽高,都是28。然后再把图片的灰度值[0, 255]
归一化到[0, 1]
,也就是除以255。
download
:True
:下载数据集;False
:不下载;默认为False
,第一次运行代码时设置为True
,之后设置为False
就行了。
数据可视化
训练数据可视化
train_data = train_file.data
train_targets = train_file.targets
print(train_data.size()) # [60000, 28, 28]
print(train_targets.size()) # [60000]
# visualization
plt.figure(figsize=(9, 9))
for i in range(9):
plt.subplot(3, 3, i+1)
plt.title(train_targets[i].numpy())
plt.axis('off')
plt.imshow(train_data[i], cmap='gray')
plt.show()
打印训练数据及标签可知:训练数据是6万张28x28的灰度图片,以及6万个标签。
测试数据可视化
test_file = datasets.MNIST(
root='./dataset/',
train=False,
transform=transforms.ToTensor()
)
test_data = test_file.data
test_targets = test_file.targets
print(test_data.size()) # [10000, 28, 28]
print(test_targets.size()) # [10000]
# visualization
plt.figure(figsize=(9, 9))
for i in range(9):
plt.subplot(3, 3, i+1)
plt.title(test_targets[i].numpy())
plt.axis('off')
plt.imshow(test_data[i], cmap='gray')
plt.show()
打印测试数据及标签可知:测试数据是1万张28x28的灰度图片,以及1万个标签。
数据加载器
train_loader = DataLoader(
dataset=train_file,
batch_size=BATCH_SIZE,
shuffle=True
)
test_loader = DataLoader(
dataset=test_file,
batch_size=BATCH_SIZE,
shuffle=False
)
dataset
:指定数据集。
batch_size
:批次的大小。
shuffle
: True:打乱数据顺序;False:不打乱数据顺序。一般训练的时候打乱顺序会取得更好的效果,测试的时候不需要打乱顺序。
模型结构
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv = nn.Sequential(
# [BATCH_SIZE, 1, 28, 28]
nn.Conv2d(1, 32, 5, 1, 2),
# [BATCH_SIZE, 32, 28, 28]
nn.ReLU(),
nn.MaxPool2d(2),
# [BATCH_SIZE, 32, 14, 14]
nn.Conv2d(32, 64, 5, 1, 2),
# [BATCH_SIZE, 64, 14, 14]
nn.ReLU(),
nn.MaxPool2d(2)
# [BATCH_SIZE, 64, 7, 7]
)
self.fc = nn.Linear(64 * 7 * 7, 10)
def forward(self, x):
x = self.conv(x)
x = x.view(x.size(0), -1)
y = self.fc(x)
return y
创建模型
model = CNN().to(device)
optim = torch.optim.Adam(model.parameters(), LR)
lossf = nn.CrossEntropyLoss()
训练模型
for epoch in range(EPOCH):
for step, (data, targets) in enumerate(train_loader):
optim.zero_grad()
data = data.to(device)
targets = targets.to(device)
output = model(data)
loss = lossf(output, targets)
loss.backward()
optim.step()
- 梯度清零
- 将图片和标签加载到
GPU
或CPU
中 - 将图片传入模型中训练
- 根据损失函数计算输出结果和标签的误差
- 误差的反向传播
- 更新参数
计算损失
loss = 0
total = 0
correct = 0
with torch.no_grad():
for data, targets in loader:
data = data.to(device)
targets = targets.to(device)
output = model(data)
loss += lossf(output, targets)
correct += (output.argmax(1) == targets).sum()
total += data.size(0)
loss = loss.item()/len(test_loader)
acc = correct.item()/total
计算模型对于整个训练集或测试集的损失和准确率。
训练过程打印
EPOCH: 01/10 STEP: 469/469 LOSS: 0.0645 ACC: 0.9807 VAL-LOSS: 0.0585 VAL-ACC: 0.9799 TOTAL-TIME: 48
EPOCH: 02/10 STEP: 469/469 LOSS: 0.0405 ACC: 0.9879 VAL-LOSS: 0.0391 VAL-ACC: 0.9870 TOTAL-TIME: 52
EPOCH: 03/10 STEP: 469/469 LOSS: 0.0297 ACC: 0.9913 VAL-LOSS: 0.0320 VAL-ACC: 0.9884 TOTAL-TIME: 46
EPOCH: 04/10 STEP: 469/469 LOSS: 0.0202 ACC: 0.9939 VAL-LOSS: 0.0271 VAL-ACC: 0.9900 TOTAL-TIME: 53
EPOCH: 05/10 STEP: 469/469 LOSS: 0.0192 ACC: 0.9941 VAL-LOSS: 0.0278 VAL-ACC: 0.9900 TOTAL-TIME: 46
EPOCH: 06/10 STEP: 469/469 LOSS: 0.0150 ACC: 0.9956 VAL-LOSS: 0.0294 VAL-ACC: 0.9897 TOTAL-TIME: 47
EPOCH: 07/10 STEP: 469/469 LOSS: 0.0114 ACC: 0.9966 VAL-LOSS: 0.0245 VAL-ACC: 0.9923 TOTAL-TIME: 54
EPOCH: 08/10 STEP: 469/469 LOSS: 0.0115 ACC: 0.9967 VAL-LOSS: 0.0269 VAL-ACC: 0.9906 TOTAL-TIME: 45
EPOCH: 09/10 STEP: 469/469 LOSS: 0.0094 ACC: 0.9972 VAL-LOSS: 0.0278 VAL-ACC: 0.9909 TOTAL-TIME: 47
EPOCH: 10/10 STEP: 469/469 LOSS: 0.0077 ACC: 0.9977 VAL-LOSS: 0.0278 VAL-ACC: 0.9916 TOTAL-TIME: 35
EPOCH
:训练的轮数
STEP
:数据加载器迭代的次数
LOSS
:整个训练集的平均损失
ACC
:整个训练集的准确率
VAL-LOSS
:整个测试集的平均损失
VAL-ACC
:整个测试集的准确率
TOTAL-TIME
:训练每轮消耗的时间(单位:秒)。
以上是在英伟达GTX 1050 Ti
训练的结果,大概每轮40
秒;如果用GTX 1080 Ti
每轮大概是20
秒。
保存最佳模型
temp = 0
if val_acc > temp:
torch.save(model.state_dict(), 'model.pt')
temp = val_acc
根据EPOCH
设置的值不同,模型会训练很多轮,为了以后方便加载,我们将最佳的模型保存下来。
首先,我们要明白:什么样的模型是最佳的?当然,对于一个好的模型来说,相比训练集,我们更看重的是它在测试集上的效果。
测试集上有两个指标可以衡量模型的好坏,损失和准确率。我们这里就保存准确率最高的。(当然损失最小的也行)
| BEST-MODEL | EPOCH: 07/10 STEP: 469/469 LOSS: 0.0114 ACC: 0.9966 VAL-LOSS: 0.0245 VAL-ACC: 0.9923
就上述10轮的训练结果来看,准确率最高的是第7轮,所以我们需要把第7轮训练的模型保存下来。
加载模型
model = CNN()
model.load_state_dict(torch.load('model.pt', map_location=torch.device('cpu')))
model.eval()
更详细的pytorch
保存和加载模型的方法可以看我翻译的这一篇官方文档。
测试模型
打开Windows
系统自带的画图应用,创建28x28
的画布,背景颜色设置为黑色,前景色为白色,然后随便写几个数字,保存为png
或其它格式的图片。
这里我手写了10个黑底白字的数字:
然后加载最佳模型,预测结果如下:(单次预测准确率100%)
pred: tensor([0, 1, 2, 3, 4, 5, 6, 7, 2, 9])
true: tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
我也手写了10个白底黑字:
测试效果很不理想,毕竟模型是在黑底白字的数据集上训练的嘛:)(单次预测准确率30%)
pred: tensor([3, 2, 2, 3, 2, 5, 5, 2, 2, 2])
true: tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
部署模型
我们不可能每次都打开画图软件写一个数字来测试吧?
pytorch
模型线上部署的教程正在制作中,敬请期待……
引用参考
https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html