跳到主要内容

PyTorch 图像分类

图像分类是计算机视觉中的一项基本任务，其目标是将输入的图像分配到预定义的类别中。PyTorch 是一个强大的深度学习框架，特别适合用于构建和训练卷积神经网络（CNN）来完成图像分类任务。在本教程中，我们将从零开始，逐步构建一个简单的图像分类模型。

1. 什么是图像分类？​

图像分类是指将输入的图像分配到预定义的类别中。例如，给定一张猫或狗的图片，模型需要判断它是猫还是狗。为了实现这一目标，我们通常使用卷积神经网络（CNN），因为它们能够有效地捕捉图像中的空间特征。

2. 准备工作​

在开始之前，确保你已经安装了 PyTorch 和 torchvision。你可以通过以下命令安装它们：

pip install torch torchvision

3. 数据准备​

图像分类的第一步是准备数据。我们将使用 torchvision 提供的 CIFAR-10 数据集，它包含了 10 个类别的 60000 张 32x32 彩色图像。

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)

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

transforms.Normalize 用于将图像的像素值归一化到 [-1, 1] 范围内，这对于训练神经网络非常重要。

4. 构建卷积神经网络​

接下来，我们将定义一个简单的卷积神经网络。这个网络包含两个卷积层和三个全连接层。

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()
提示

nn.Conv2d 是卷积层，nn.MaxPool2d 是最大池化层，nn.Linear 是全连接层。F.relu 是激活函数，用于引入非线性。

5. 定义损失函数和优化器​

为了训练模型，我们需要定义损失函数和优化器。这里我们使用交叉熵损失函数和随机梯度下降（SGD）优化器。

import torch.optim as optim

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

6. 训练模型​

现在我们可以开始训练模型了。我们将遍历数据集多次（称为 epoch），并在每个 epoch 中更新模型的参数。

for epoch in range(2):  # 遍历数据集两次

    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:  # 每 2000 个小批量打印一次
            print(f'[Epoch {epoch + 1}, Batch {i + 1}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

print('Finished Training')

7. 测试模型​

训练完成后，我们需要在测试集上评估模型的性能。

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(f'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%')

8. 实际应用场景​

图像分类在许多实际应用中都有广泛的应用，例如：

  • 医学影像分析：自动识别医学图像中的病变区域。
  • 自动驾驶：识别道路上的行人、车辆和交通标志。
  • 安防监控：识别监控视频中的异常行为。

9. 总结​

在本教程中，我们学习了如何使用 PyTorch 构建和训练一个简单的卷积神经网络来进行图像分类。我们从数据准备开始，逐步构建了模型，并进行了训练和测试。希望你能通过本教程掌握图像分类的基本流程，并能够将其应用到实际问题中。

10. 附加资源与练习​

警告

在实际应用中，模型的性能可能会受到数据质量、模型复杂度等因素的影响。因此，建议在实际项目中仔细调优模型。