PyTorch 图像分类教程

图像分类是计算机视觉领域的一个重要任务，而 PyTorch 是当前最受欢迎的深度学习框架之一。本教程将带你一步步了解如何在 PyTorch 中进行图像分类。

基础知识

在开始之前，请确保你已经安装了 PyTorch。如果没有安装，可以访问 PyTorch 官网了解如何安装。

环境搭建

首先，我们需要搭建一个合适的环境来进行图像分类实验。以下是一个简单的步骤：

安装 PyTorch。
准备数据集，例如 CIFAR-10。
构建模型。
训练模型。
评估模型。

数据集

在图像分类任务中，数据集的选择至关重要。以下是一些常用的图像分类数据集：

CIFAR-10
ImageNet
MNIST

CIFAR-10

CIFAR-10 是一个包含 10 个类别、共 10 万张 32x32 图像的数据集。每个类别有 1 万张训练图像和 1 万张测试图像。

模型构建

在 PyTorch 中，我们可以使用 torch.nn 模块来构建模型。以下是一个简单的卷积神经网络示例：

import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.max_pool2d(x, 2)
        x = nn.functional.relu(self.conv2(x))
        x = nn.functional.max_pool2d(x, 2)
        x = x.view(-1, 64 * 8 * 8)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

训练模型

在构建好模型后，我们需要进行训练。以下是一个简单的训练流程：

定义损失函数和优化器。
训练模型。
保存模型。

import torch.optim as optim

model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(10):  # 训练 10 个 epoch
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

评估模型

训练完成后，我们需要评估模型在测试集上的表现。以下是一个简单的评估流程：

加载测试集。
测试模型。
输出准确率。

model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for data, target in test_loader:
        outputs = model(data)
        _, predicted = torch.max(outputs.data, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()

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

总结

通过本教程，你了解了如何在 PyTorch 中进行图像分类。希望这个教程对你有所帮助。如果你想要了解更多关于 PyTorch 的知识，可以访问我们的 PyTorch 教程页面。