PyTorch项目笔记（二）CIFAR-10图像分类

最新推荐文章于 2024-03-20 09:38:47 发布

Xyzz1223

最新推荐文章于 2024-03-20 09:38:47 发布

阅读量2k

点赞数 3

分类专栏： PyTorch 文章标签： pytorch 分类 python

本文链接：https://blog.csdn.net/a1016249126/article/details/129643589

版权

PyTorch 专栏收录该内容

3 篇文章 1 订阅

订阅专栏

本文将在CIFAR10的基础上介绍如何处理数据、图像分类以及使用GPU加速模型训练。

1 CIFAR-10数据集介绍

CIFAR-10数据集共有60000张尺寸为3 × 32 × 32的图像，包含以下10个类别：

classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

其中训练集图像50000张，测试集图像10000张。

2 如何加载数据集

我们使用Pytorch提供的视觉工具包torchvision加载CIFAR-10数据集。在加载数据集时，我们可以对训练集和测试集分别进行预处理操作，此处需注意只有训练集需要进行图像增强。此外，使用torchvision包可以快速构建出DataLoader对象。具体代码如下：

#导入所需的包
import torch
import torchvision
import torchvision.transforms as transforms

#训练集图像预处理：
#对测试集进行图像增强可缓解过拟合的情况
train_transform = transforms.Compose([transforms.RandomRotation(degrees=10), #随机旋转
                                      transforms.RandomHorizontalFlip(p=0.5), #随机水平翻转
                                      transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1), #改变图像属性
                                      transforms.ToTensor(), #转Tensor
                                      transforms.Normalize( 
                                        mean=[0.5, 0.5, 0.5],
                                        std=[0.5, 0.5, 0.5]) #归一化
                                     ])

# 测试集图像预处理
test_transform = transforms.Compose([transforms.ToTensor(), #转Tensor
                                     transforms.Normalize(
                                        mean=[0.5, 0.5, 0.5],
                                        std=[0.5, 0.5, 0.5]) #归一化
                                    ])

#加载训练集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=train_transform)
#为训练集创建DataLoader对象
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=test_transform)
#为测试集创建DataLoader对象
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)
#定义数据集的类别
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

下面我们将使用matplotlib包显示数据集中的部分图像，具体代码如下：

#导入所需的包
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 = dataiter.next()

#显示图像
imshow(torchvision.utils.make_grid(images))
#打印对应图像的标签
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

运行结果如下图所示。

请添加图片描述

 frog  ship  deer  dog

3 定义CNN网络模型

此处定义一个简单的前馈卷积神经网络，具体代码如下：

#导入所需的包
import torch.nn as nn
import torch.nn.functional as F
#定义Net类
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

4 使用GPU加速

与Tensor转移到GPU上计算一样，我们只需把神经网络模块转移到GPU上即可，具体代码如下：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

若输出以下结果，则说明成功将神经网络模块转移到CUDA设备上。

cuda:0

需要注意的是，当网络移动到CUDA设备后，输入到网络中的张量也需要先移动到CUDA设备上。因为PyTorch只能在同一个设备上做矩阵操作，具体代码如下：

net.to(device)

inputs, labels = inputs.to(device), labels.to(device)

5 训练网络

在训练网络之前，我们需要先定义损失函数和优化器，这里使用交叉熵（CrossEntropy）作为损失函数，支持动量的SGD作为优化器，具体代码如下：

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

训练模型时只需要在数据迭代器上循环传给网络和优化器的输入，具体代码如下：

for epoch in range(epoches):
    running_loss = 0.0
    count = 0
    startTime = time.time()
    for i, data in enumerate(trainloader, 0):
        # get the inputs
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = classifyNet(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        count = count + 1
	totalTime = time.time() - startTime
    print('epoch: %d loss: %.3f time: %.3f' %(epoch + 1, running_loss / count, totalTime))

print('Finished Training')

我们以5个epoch为例，训练结果如下图所示。
请添加图片描述

6 在测试集上评估

首先，我们对整个测试集中断所有图像进行预测，并计算整体的准确率，具体代码如下：

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = classifyNet(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))

运行结果如下：

Accuracy of the network on the 10000 test images: 59 %

如果模型是随机预测的话，那么准确率应该是10%，而我们的训练两个epoch得到的模型，准确率为59%。这说明网络还是学到了一些东西。为了进行精细化分析，下面我们看看模型在每一个类别上的准确率，具体代码如下：

class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = classifyNet(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1


for i in range(10):
    print('Accuracy of %5s : %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))

运行结果如下所示：

Accuracy of plane : 72 %
Accuracy of   car : 87 %
Accuracy of  bird : 31 %
Accuracy of   cat : 40 %
Accuracy of  deer : 57 %
Accuracy of   dog : 45 %
Accuracy of  frog : 77 %
Accuracy of horse : 69 %
Accuracy of  ship : 49 %
Accuracy of truck : 60 %