利用pytorch实现对CIFAR-10数据集的分类

在机器学习领域,图像分类是一项非常重要的任务。CIFAR-10是一个经典的图像分类数据集,其中包含10个不同的类别,每个类别有6000张32x32像素的图像。对于初学者来说,使用PyTorch实现CIFAR-10数据集的分类非常有帮助,因为它是一项非常具有挑战性而且有许多变量和参数需要掌握的任务。

1. CIFAR-10数据集的加载与预处理

在PyTorch中使用CIFAR-10数据集非常方便,我们可以使用torchvision.datasets.CIFAR10函数直接加载数据。在加载数据之前,我们需要对数据进行预处理,包括将数据转换为Tensor类型,将数据进行标准化处理等。

1.1 数据集的加载

代码如下:

import torch

import torchvision

import torchvision.transforms as transforms

transform = transforms.Compose(

[transforms.ToTensor(),

transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,

download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,

shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,

download=True, transform=transform)

testloader = torch.utils.data.DataLoader(testset, batch_size=4,

shuffle=False, num_workers=2)

在上述代码中,我们首先定义了一个transform的变量,这个变量包含了两个预处理步骤:将数据转换为tensor类型,然后将数据进行标准化处理。在加载数据集时,我们需要将这个transform对象传递给CIFAR-10函数。

接着,我们定义了trainset和testset两个变量来保存训练集和测试集。我们可以使用train=True和train=False来指定是加载训练集还是测试集。在创建数据加载器时,我们还需要指定batch_size,这指定了每个批次包含的图像数量。shuffle参数表示是否要对数据进行打乱顺序,num_workers参数表示使用多少个线程来加载数据。

1.2 数据展示

接下来,我们可以通过以下代码来展示数据集中的部分图片:

import matplotlib.pyplot as plt

import numpy as np

# 定义标签

classes = ('plane', 'car', 'bird', 'cat',

'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# 随机获取一些训练数据

dataiter = iter(trainloader)

images, labels = dataiter.next()

# 绘制图像

def imshow(img):

img = img / 2 + 0.5 # 将图像的像素值标准化到 [0, 1] 区间

npimg = img.numpy()

plt.imshow(np.transpose(npimg, (1, 2, 0)))

plt.show()

# 展示图像

imshow(torchvision.utils.make_grid(images))

# 展示类别标签

print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

在上述代码中,我们首先定义了类别标签,然后使用iter函数和next函数来随机获取一些训练数据。然后,我们使用imshow函数来绘制图像,使用print语句来展示类别标签。运行以上代码,会输出4张随机选择的图片以及它们对应的标签。

2. 模型的定义和训练

定义和训练模型是图像分类任务中最重要的部分。在PyTorch中,模型可以通过继承torch.nn.Module类来定义,我们只需要实现__init__函数和forward函数即可。在训练模型时,我们需要定义损失函数和优化器。

2.1 模型的定义

在本文中,我们使用一个非常简单的卷积神经网络(Convolutional Neural Network,CNN)来对CIFAR-10数据集进行分类。CNN是一种特殊的神经网络,它可以有效地处理图像数据。

代码如下:

import torch.nn as nn

import torch.nn.functional as F

class Net(nn.Module):

def __init__(self):

super(Net, self).__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 = x.view(-1, 16 * 5 * 5)

x = F.relu(self.fc1(x))

x = F.relu(self.fc2(x))

x = self.fc3(x)

return x

net = Net()

在上述代码中,我们定义了一个名为Net的类。这个类继承自nn.Module类,并实现了__init__函数和forward函数。在__init__函数中,我们定义了两个卷积层、两个全连接层和一个池化层。在forward函数中,我们使用卷积层、池化层和全连接层来构建神经网络。

2.2 训练模型

接下来,我们可以通过以下代码来训练模型:

import torch.optim as optim

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):

# 获取输入数据

inputs, labels = data

# 将梯度重置为0

optimizer.zero_grad()

# 正向传递、反向传递、优化

outputs = net(inputs)

loss = criterion(outputs, labels)

loss.backward()

optimizer.step()

# 打印统计信息

running_loss += loss.item()

if i % 2000 == 1999: # 每2000个批次打印一次

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

(epoch + 1, i + 1, running_loss / 2000))

running_loss = 0.0

print('Finished Training')

在上述代码中,我们首先定义了损失函数和优化器。使用nn.CrossEntropyLoss作为损失函数,使用optim.SGD作为优化器。

然后,我们对训练集进行循环遍历,每次从训练集中获取一个批次的数据。在正向传播时,我们将批次数据作为输入,通过神经网络,得到预测值。在反向传播时,我们计算损失值和梯度,并使用optimizer.step()函数来更新神经网络中的参数。

在循环中,我们还定义了一个running_loss变量来保存损失值,每2000个批次打印一次损失值。

3. 模型的评估和测试

在训练模型之后,我们需要对模型进行评估和测试。评估模型时,我们通常使用准确率来衡量模型的性能。在测试模型时,我们需要通过模型来预测一组数据,并根据预测结果来评估模型的性能。

3.1 模型的评估

评估模型的准确率是非常简单的,我们只需要对测试集中的每个图像进行预测,并将预测结果与实际结果进行比较。代码如下:

correct = 0

total = 0

with torch.no_grad():

for data in testloader:

# 获取输入数据

images, labels = data

# 使用模型进行预测

outputs = net(images)

_, predicted = torch.max(outputs.data, 1)

# 统计预测结果

total += labels.size(0)

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

print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))

在上述代码中,我们首先将模型设置为评估模式。然后,我们对测试集进行循环遍历,对每个图像进行预测。在预测中,我们使用torch.max函数来获取预测结果的索引值,并与实际结果进行比较。在循环中,我们还定义了一个correct变量来保存预测正确的图像数量,一个total变量来保存总图像数量。

3.2 模型的测试

接下来,我们可以使用以下代码来测试模型:

# 获取随机图像数据

dataiter = iter(testloader)

images, labels = dataiter.next()

# 显示图片

imshow(torchvision.utils.make_grid(images))

print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))

# 对图像进行预测

outputs = net(images)

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

print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]

for j in range(4)))

在上述代码中,我们首先获取一些测试数据,并对这些数据进行预测。然后,我们使用imshow函数来显示这些测试数据图像,并使用print语句来输出实际类别和预测类别。

使用以上代码可得到的模型准确率为49%。由于本篇文章只是对CIFAR-10数据集分类的入门教学,因此本篇文章实现的模型并不是非常复杂,并且只利用了非常基础的算法。如果您需要实现更加高效和复杂的算法,可以在实现上述代码之后进行自己的尝试。同时,可以通过调整模型的结构、超参数等参数来提高模型的准确率。

后端开发标签