实战教程：用 PyTorch 构建你的第一个图像分类模型，从零开始！

在深度学习领域，图像分类是最基础且应用广泛的任务之一，从自动驾驶中的路标识别，到医疗影像中的疾病诊断，都离不开图像分类模型的支持。PyTorch 作为一款简洁且高效的深度学习框架，以其动态计算图和 Pythonic 的编程风格，深受开发者喜爱。本文将带领你从零开始，使用 PyTorch 构建一个属于自己的图像分类模型，让你轻松入门深度学习图像分类领域。

一、环境搭建

在开始构建模型之前，首先需要搭建好 PyTorch 运行环境。

1.1 安装 Python

PyTorch 支持 Python 3.7 及以上版本，你可以从Python 官方网站下载对应系统的 Python 安装包进行安装。安装过程中记得勾选 “Add Python to PATH” 选项，方便后续在命令行中使用 Python。

1.2 安装 PyTorch

PyTorch 的安装方式有多种，这里推荐使用 conda 或 pip 进行安装。

• 使用 conda 安装：首先确保你已经安装了 Anaconda 或 Miniconda。打开 Anaconda Prompt（Windows）或终端（Linux/macOS），根据你的 CUDA 版本选择对应的命令安装。
  • 如果你有 CUDA 11.8 环境，执行以下命令：

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

• 如果你没有 GPU，想使用 CPU 版本，执行：

conda install pytorch torchvision torchaudio cpuonly -c pytorch

• 使用 pip 安装：同样打开命令行工具，根据 CUDA 版本安装。例如，CUDA 11.8 版本：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

CPU 版本：

pip install torch torchvision torchaudio

安装完成后，可以在 Python 交互环境中输入以下代码验证 PyTorch 是否安装成功：

import torch
print(torch.__version__)

如果输出版本号，则说明安装成功。

二、数据集准备

本次教程使用经典的 CIFAR-10 数据集，该数据集包含 10 个类别，共 60000 张彩色图像，其中 50000 张用于训练，10000 张用于测试。在 PyTorch 中，可以通过torchvision库方便地下载和加载 CIFAR-10 数据集。

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

上述代码中，首先定义了数据预处理操作，将图像转换为张量并进行归一化，这有助于模型更快地收敛。然后分别下载了 CIFAR-10 的训练集和测试集，并使用DataLoader将数据封装成批次，方便后续训练和测试。

三、构建图像分类模型

这里我们构建一个简单的卷积神经网络（Convolutional Neural Network，CNN）来进行图像分类。CNN 在处理图像数据时具有天然的优势，能够自动提取图像的特征。

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 第一个卷积层，输入通道为3（彩色图像），输出通道为6，卷积核大小为5x5
        self.conv1 = nn.Conv2d(3, 6, 5)
        # 最大池化层，池化核大小为2x2
        self.pool = nn.MaxPool2d(2, 2)
        # 第二个卷积层，输入通道为6，输出通道为16，卷积核大小为5x5
        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，这是 PyTorch 中所有神经网络模型的基类。在类的构造函数中，定义了网络的各个层，包括卷积层、池化层和全连接层。forward方法定义了数据在网络中的前向传播路径，通过一系列的卷积、池化和全连接操作，最终输出预测的类别得分。

四、训练模型

模型构建完成后，需要对其进行训练，让模型学习数据中的特征和规律。训练过程中需要定义损失函数和优化器。

import torch.optim as optim

# 定义损失函数为交叉熵损失函数
criterion = nn.CrossEntropyLoss()
# 定义优化器为随机梯度下降（SGD），学习率为0.001，动量为0.9
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(10):  # 训练10个epoch
    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个mini-batch打印一次
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')

上述代码中，选择交叉熵损失函数作为衡量模型预测结果与真实标签差异的指标，使用随机梯度下降优化器来更新模型的参数。在训练循环中，每个批次的数据经过前向传播得到预测结果，计算损失后进行反向传播计算梯度，最后通过优化器更新模型参数。

五、测试模型

训练完成后，需要在测试集上评估模型的性能，查看模型在未见过的数据上的表现。

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

上述代码中，在测试过程中，关闭梯度计算，因为不需要在测试时更新模型参数。通过比较模型预测结果和真实标签，计算出模型在测试集上的准确率，以此评估模型的性能。

六、完整代码

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


# 数据预处理
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')


# 定义神经网络模型
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()

# 定义损失函数和优化器
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

print('Finished Training')

# 测试模型
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))

七、总结

通过以上步骤，我们从零开始使用 PyTorch 完成了一个图像分类模型的搭建、训练和测试。虽然这个模型相对简单，但它为我们学习深度学习图像分类任务奠定了坚实的基础。在实际应用中，你可以尝试调整模型结构、优化超参数、使用更复杂的数据集等，进一步提升模型的性能。希望本教程能帮助你顺利踏入深度学习图像分类的大门，开启更多有趣的探索之旅。