基于 PyTorch 的 CNN 图像分类器实现。


简介

由于卷积神经网络 (CNN) 能够识别图像中的空间特征,因此已广泛应用于图像分类应用中。PyTorch 是一个流行的开源机器学习库,它提供了构建和训练神经网络(包括 CNN)的支持。本文将介绍如何使用 PyTorch 构建基于 CNN 的图像分类器。

数据集

在深入探讨实现细节之前,让我们先了解一下数据集。本教程将使用 CIFAR-10 数据集,它包含 60,000 张 32x32 的彩色图像,分为 10 个类别,每个类别包含 6,000 张图像。这些类别包括卡车、船、青蛙、马、鸟、猫、鹿、汽车、飞机和汽车。数据集包含 10,000 张测试图像和 50,000 张训练图像。

数据预处理

构建基于 CNN 的图像分类器的第一步是预处理数据。在代码中,我们将对图像像素值进行归一化,使其范围在 0 到 1 之间。我们还将使用各种数据增强技术来扩展数据集并减少过拟合。数据增强技术包括随机旋转、随机裁剪和随机水平翻转。

构建 CNN

数据预处理完成后,我们可以开始构建 CNN。我们的 CNN 架构将包含多个卷积层、最大池化层和全连接层。

训练 CNN

定义 CNN 后,我们就可以使用 CIFAR-10 数据集来训练它。为此,我们将利用 PyTorch 内置的神经网络训练功能,包括定义损失函数和优化器。

我们将使用交叉熵损失函数,它通常用于分类任务。我们将使用随机梯度下降 (SGD) 作为优化器,学习率为 0.001,动量为 0.9。

评估 CNN

成功训练 CNN 后,我们可以通过分析测试数据集来评估其性能。我们可以通过迭代数据集并评估 CNN 的准确率来实现这一点。

代码

以下是用 PyTorch 实现的简单基于 CNN 的图像分类器的示例,

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim

# Define the CNN architecture
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(nn.functional.relu(self.conv1(x)))
      x = self.pool(nn.functional.relu(self.conv2(x)))
      x = x.view(-1, 16 * 5 * 5)
      x = nn.functional.relu(self.fc1(x))
      x = nn.functional.relu(self.fc2(x))
      x = self.fc3(x)
      return x

# Preprocess the data
transform = transforms.Compose(
   [transforms.RandomHorizontalFlip(),
   transforms.RandomCrop(32, padding=4),
   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=32, 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=32, shuffle=False, num_workers=2)

# Train the network
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(10):
   running_loss = 0.0
   for i, data in enumerate(trainloader, 0):
      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('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
         running_loss = 0.0

# Evaluate the network
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))

输出

[1,  2000] loss: 2.168
[1,  4000] loss: 1.828
[1,  6000] loss: 1.639
[1,  8000] loss: 1.552
[1, 10000] loss: 1.490
[2,  2000] loss: 1.407
[2,  4000] loss: 1.378
[2,  6000] loss: 1.324
[2,  8000] loss: 1.297
[2, 10000] loss: 1.269
[3,  2000] loss: 1.190
[3,  4000] loss: 1.170
[3,  6000] loss: 1.142
[3,  8000] loss: 1.126
[3, 10000] loss: 1.108
[4,  2000] loss: 1.031
[4,  4000] loss: 1.037
[4,  6000] loss: 1.016
[4,  8000] loss: 1.005
[4, 10000] loss: 1.002
[5,  2000] loss: 0.932
[5,  4000] loss: 0.944
[5,  6000] loss: 0.932
[5,  8000] loss: 0.913
[5, 10000] loss: 0.913
[6,  2000] loss: 0.831
[6,  4000] loss: 0.835
[6,  6000] loss: 0.846
[6,  8000] loss: 0.835
[6, 10000] loss: 0.829
[7,  2000] loss: 0.745
[7,  4000] loss: 0.759
[7,  6000] loss: 0.762
[7,  8000] loss: 0.763
[7, 10000] loss: 0.753
[8,  2000] loss: 0.670
[8,  4000] loss: 0.696
[8,  6000] loss: 0.690
[8,  8000] loss: 0.701
[8, 10000] loss: 0.690
[9,  2000] loss: 0.614
[9,  4000] loss: 0.635
[9,  6000] loss: 0.646
[9,  8000] loss: 0.641
[9, 10000] loss: 0.631
[10,  2000] loss: 0.562
[10,  4000] loss: 0.577
[10,  6000] loss: 0.597
[10,  8000] loss: 0.592
[10, 10000] loss: 0.586
Accuracy of the network on the 10000 test images: 68 %

此输出显示了训练后网络在测试集上的准确率以及每次训练迭代时的损失。如您所见,测试集上的准确率峰值为 68%,并且损失在 10 次训练时期内逐渐下降。此结果可能不是最先进的,但它确实展示了使用 PyTorch 训练 CNN 进行图像分类的基本步骤。

结论

总之,本文介绍了如何使用 PyTorch 构建基于 CNN 的图像分类器。该过程包括数据预处理,包括归一化和数据增强。CNN 架构由卷积层、最大池化层和全连接层组成。训练使用交叉熵损失和随机梯度下降优化器在 CIFAR-10 数据集上进行。然后在测试数据集上评估和评估 CNN 的准确率。有了这些知识,您现在就可以开始构建自己的图像分类模型并进一步了解计算机视觉。

更新于: 2023年7月12日

126 次浏览

开启你的 职业生涯

通过完成课程获得认证

开始学习
广告