训练分类器
一小时速通pytorch系列文章:
一小时速通Pytorch之自动梯度(Autograd)和计算图(Computational Graph)(二)(完结)
数据集在此:
通过网盘分享的文件:data.rar
链接: https://pan.baidu.com/s/18FCDSwwUCyg4BpTZ73dbLg?pwd=xixi 提取码: xixi
目前为止,我们已经学习了如何定义一个神经网络、计算损失以及更新网络的权重参数。但是我们必须要思考一个问题:
数据在哪?
一般地,当我们去处理图像 ,文本 ,音视频 数据的时候,我们都可以使用一些python标准库把这些数据转换成一个numpy数组,然后再转换成torch.*Tensor类型
- 对于图像数据,我们可以采用**
Pillow,OpenCV**这两个库 - 对于音频数据,可以采用**
scipy,librosa**这两个库 - 对于文本数据,无论是原始Python还是基于**
Cython的加载,或者NLTK和SpaCy**都是可以的
针对视觉 这个领域的数据,pytorch开发了torchvison这个包,这个包的torchvision.datasets部分包含了ImageNet,CIFAR10,MINIST等一些常用的数据集;torch.utils.data.DataLoader则是一个用于图像的数据加载器。
pytorch可以让我们方便的导入图像数据,避免写一些模板性的代码。
在这个教程中,我们将使用CIFAR10这个数据集。这个数据的标签(分类)数据有airplane,automobile,bird,cat,deer,dog,frog,horse,ship,truck10类。数据集中的数据(图片)都是32 * 32 * 3,即分辨率是 32 * 32的3通道图片。
数据集示例:

训练一个图像分类器
接下来我们从这几个步骤展开:
- 使用
torchvision从CIFAR10中加载和归一化数据,训练以及测试。 - 定义一个卷积神经网络
- 定义一个损失函数
- 在训练集上训练网络
- 用验证机来测试网络的效果
1. 加载和归一化CIFAR10数据集
注意事项:
在Windows和MacOs下导入数据集的加载器的时候,需要指定num_workers=0避免多线程错误
示例:
python
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0)
我们通过torchvision可以很方便导入CIFAR10数据集
示例:
python
import torch
import torchvision
from torchvision import transforms
torchvision.datasets数据集输出的都是PILImage类型的图像,并且范围在[0,1]之间,我们把这些数据转换成Tensor并且数据范围在[-1,1]之间
示例:
python
import torch
import torchvision
from torchvision import transforms
# 调整数据的范围
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
batch_size = 4
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=0)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
示例输出:
0%| | 0.00/170M [00:00<?, ?B/s]
0%| | 459k/170M [00:00<00:37, 4.48MB/s]
3%|▎ | 4.75M/170M [00:00<00:06, 26.8MB/s]
5%|▌ | 9.37M/170M [00:00<00:04, 35.5MB/s]
8%|▊ | 14.3M/170M [00:00<00:03, 40.6MB/s]
11%|█ | 19.1M/170M [00:00<00:03, 43.3MB/s]
14%|█▍ | 24.3M/170M [00:00<00:03, 46.1MB/s]
17%|█▋ | 29.5M/170M [00:00<00:02, 47.7MB/s]
20%|██ | 34.3M/170M [00:00<00:02, 46.8MB/s]
23%|██▎ | 39.1M/170M [00:00<00:02, 47.0MB/s]
26%|██▌ | 43.9M/170M [00:01<00:02, 47.4MB/s]
29%|██▊ | 48.7M/170M [00:01<00:02, 45.4MB/s]
31%|███▏ | 53.3M/170M [00:01<00:02, 45.4MB/s]
34%|███▍ | 57.9M/170M [00:01<00:02, 44.7MB/s]
37%|███▋ | 62.4M/170M [00:01<00:02, 42.7MB/s]
39%|███▉ | 66.7M/170M [00:01<00:02, 43.0MB/s]
42%|████▏ | 71.6M/170M [00:01<00:02, 44.5MB/s]
45%|████▍ | 76.1M/170M [00:01<00:02, 43.3MB/s]
47%|████▋ | 80.4M/170M [00:01<00:02, 43.3MB/s]
50%|████▉ | 84.8M/170M [00:01<00:01, 43.1MB/s]
52%|█████▏ | 89.1M/170M [00:02<00:01, 42.6MB/s]
55%|█████▍ | 93.4M/170M [00:02<00:01, 42.4MB/s]
57%|█████▋ | 97.7M/170M [00:02<00:01, 42.4MB/s]
60%|█████▉ | 102M/170M [00:02<00:01, 41.9MB/s]
62%|██████▏ | 106M/170M [00:02<00:01, 41.3MB/s]
65%|██████▍ | 110M/170M [00:02<00:01, 40.5MB/s]
67%|██████▋ | 114M/170M [00:02<00:01, 38.6MB/s]
69%|██████▉ | 118M/170M [00:02<00:01, 38.0MB/s]
72%|███████▏ | 123M/170M [00:02<00:01, 39.3MB/s]
74%|███████▍ | 127M/170M [00:03<00:01, 40.0MB/s]
77%|███████▋ | 131M/170M [00:03<00:00, 40.0MB/s]
79%|███████▉ | 135M/170M [00:03<00:00, 39.7MB/s]
82%|████████▏ | 139M/170M [00:03<00:00, 41.0MB/s]
84%|████████▍ | 143M/170M [00:03<00:00, 40.5MB/s]
86%|████████▋ | 147M/170M [00:03<00:00, 40.2MB/s]
89%|████████▉ | 151M/170M [00:03<00:00, 39.7MB/s]
91%|█████████▏| 156M/170M [00:03<00:00, 40.0MB/s]
94%|█████████▎| 160M/170M [00:03<00:00, 40.4MB/s]
96%|█████████▌| 164M/170M [00:03<00:00, 39.4MB/s]
98%|█████████▊| 168M/170M [00:04<00:00, 39.5MB/s]
100%|██████████| 170M/170M [00:04<00:00, 41.4MB/s]
网络"好"的读者可以采取上文从网络上下载数据集来进行加载。网络"不好"的读者,先下载资源中的data压缩包然后解压放在运行python脚本运行的同级目录下即可。
示例,网络"不好"的导入数据集的方法
python
import torch
import torchvision
from torchvision import transforms
# 调整数据的范围
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
batch_size = 4
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=0)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
接下来,我们定义一个函数来展示数据集中的一些数据
示例:
python
import matplotlib.pyplot as plt
import numpy as np
def imshow(img):
img = img / 2 + 0.5 # 反归一化
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
# 从训练集中获取一些随机的数据
dataiter = iter(trainloader)
images, labels = next(dataiter)
# 展示图像
imshow(torchvision.utils.make_grid(images))
# 打印标签
print("图中的类别分别是:")
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))
示例输出:

bird ship dog plane
2. 定义一个卷积神经网络
从之前第三节定义的神经网络章节复制神经网络的定义部分,并把之前定义的单通道的信息改成三通道。具体表现在第一层的卷积层的输入从1 变成了3
示例:
python
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
net = Net()
3. 定义一个损失函数和优化器
这里采用交叉熵损失函数(在多分类问题中常用)和使用随机梯度下降方法的优化器。
示例:
python
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
4. 训练网络
这里,我们只需遍历数据迭代器,将输入馈送到网络并进行优化。
示例:
python
for epoch in range(2): # 训练集完整迭代的次数
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
optimizer.zero_grad()
# 前向传播,反向传播,优化器
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 打印统计数据
running_loss += loss.item()
if i % 2000 == 1999:
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
running_loss = 0.0
print('训练结束')
# 保存参数
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)
示例输出:
[1, 2000] loss: 2.289
[1, 4000] loss: 2.007
[1, 6000] loss: 1.743
[1, 8000] loss: 1.606
[1, 10000] loss: 1.526
[1, 12000] loss: 1.477
[2, 2000] loss: 1.395
[2, 4000] loss: 1.389
[2, 6000] loss: 1.360
[2, 8000] loss: 1.310
[2, 10000] loss: 1.321
[2, 12000] loss: 1.305
训练结束
5. 在验证集上测试数据
我们以及在训练集上训练了两次。但是我们要检查一下网络是否学习到任何信息。
我们将通过预测神经网络输出的类标签来检查这一点,并将其与基本事实进行比较。如果预测正确,我们将样本添加到正确预测列表中。
首先,让我们从验证集中显示一张图像来让我们有一个直观的了解。
示例:
python
dataiter = iter(testloader)
images, labels = next(dataiter)
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))
示例输出:

GroundTruth: cat ship ship plane
接下来,让我们重新加载保存的模型(注意:这里不需要保存和重新加载模型,我们这样做只是为了说明如何操作)
示例:
python
net = Net()
PATH = './cifar_net.pth'
net.load_state_dict(torch.load(PATH, weights_only=True))
让我们把验证集传进去,看一下神经网络的推断能力如何
示例:
python
outputs = net(images)
outputs是一个十个类别的数据,某个类别的得分越高,神经网络就把该张照片推断成这个类别。因此我们需要获取最高得分类别的数据。
示例:
python
rint('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}' for j in range(4)))
示例输出:
GroundTruth: cat ship ship plane
Predicted: cat ship ship ship
单从这几张照片来看,结果还算准确,让我们看一下在整个数据集上这个网络的表现如何:
示例:
python
correct = 0
total = 0
# 验证阶段不需要梯度信息,暂时关闭
with torch.no_grad():
for data in testloader:
images, labels = data
# 计算模型输出值
outputs = net(images)
# 我们选择得分最高的类作为预测
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'在10000张数据上的准确率为: {100 * correct // total} %')
示例输出:
在10000张数据上的准确率为: 55 %
看来经过两次迭代效果比随机值(十分之一的概率)好了一点,看起来模型确实学到了一点东西。但是具体是哪几个类学的比较好,哪几个类学的比较差呢?让我们再统计一下:
示例:
python
# 准备为每个类别统计预测
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}
# 这里也是不需要梯度信息
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = net(images)
_, predictions = torch.max(outputs, 1)
# 收集每个类别的正确预测
for label, prediction in zip(labels, predictions):
if label == prediction:
correct_pred[classes[label]] += 1
total_pred[classes[label]] += 1
# 打印每一类的准确度
for classname, correct_count in correct_pred.items():
accuracy = 100 * float(correct_count) / total_pred[classname]
print(f'类别: {classname:5s} 的准确率为: {accuracy:.1f} %')
示例输出:
类别: plane 的准确率为: 61.9 %
类别: car 的准确率为: 80.2 %
类别: bird 的准确率为: 54.7 %
类别: cat 的准确率为: 28.8 %
类别: deer 的准确率为: 37.4 %
类别: dog 的准确率为: 45.0 %
类别: frog 的准确率为: 70.9 %
类别: horse 的准确率为: 64.8 %
类别: ship 的准确率为: 64.1 %
类别: truck 的准确率为: 44.9 %
我们上述只是在CPU上进行训练,接下来把数据搬到GPU再进行训练
在GPU上进行训练
正如我们把张量搬运到GPU上一样,我们的神经网络也可以搬运到GPU上。
如果我们有可用的cuda设备,让我们将我们的设备定义为第一个可见的cuda设备:
python
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
示例输出:
cuda:0
这里笔者是有一台可以CUDA的设备,接下来的内容都是假设我们的device是cuda设备,把我们之前方法涉及到的参数和缓冲区的数据都搬运到cuda上。
示例:
python
net.to(device)
inputs, labels = data[0].to(device), data[1].to(device)
这里我们用GPU训练了一下网络感觉和CPU没有多大差别,这是因为我们的网络很小,参数也不大,在CPU和GPU上差距不明显。
至此,训练分类器一章完结,如果想要参考多GPU并行训练,可以参考这里
完整代码
python
import torch
import torchvision
from torchvision import transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
# 调整数据的范围
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
batch_size = 4
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=0)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# 显示图像的函数
def imshow(img):
img = img / 2 + 0.5 # 反归一化
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
# # 从训练集中获取一些随机的数据
# dataiter = iter(trainloader)
# images, labels = next(dataiter)
#
# # 展示图像
# imshow(torchvision.utils.make_grid(images))
# # 打印标签
# print("图中的类别分别是:")
# print(' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
def train():
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
for epoch in range(2): # 训练集完整迭代的次数
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
optimizer.zero_grad()
# 前向传播,反向传播,优化器
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 打印统计数据
running_loss += loss.item()
if i % 2000 == 1999:
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
running_loss = 0.0
print('训练结束')
# 保存参数
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)
def test():
dataiter = iter(testloader)
images, labels = next(dataiter)
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))
net = Net()
PATH = './cifar_net.pth'
net.load_state_dict(torch.load(PATH, weights_only=True))
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}' for j in range(4)))
correct = 0
total = 0
# 验证阶段不需要梯度信息,暂时关闭
with torch.no_grad():
for data in testloader:
images, labels = data
# 计算模型输出值
outputs = net(images)
# 我们选择得分最高的类作为预测
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'在10000张数据上的准确率为: {100 * correct // total} %')
# 准备为每个类别统计预测
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}
# 这里也是不需要梯度信息
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = net(images)
_, predictions = torch.max(outputs, 1)
# 收集每个类别的正确预测
for label, prediction in zip(labels, predictions):
if label == prediction:
correct_pred[classes[label]] += 1
total_pred[classes[label]] += 1
# 打印每一类的准确度
for classname, correct_count in correct_pred.items():
accuracy = 100 * float(correct_count) / total_pred[classname]
print(f'类别: {classname:5s} 的准确率为: {accuracy:.1f} %')
if __name__ == '__main__':
# train()
# test()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)