从零开始：使用PyTorch实现图像分类（附完整代码与注释）

引言

图像分类是计算机视觉领域的基础任务，广泛应用于人脸识别、医学影像分析、自动驾驶等场景。PyTorch作为主流深度学习框架，以其动态计算图和易用性深受开发者喜爱。本文将通过完整代码和详细注释，展示如何使用PyTorch实现一个基础的图像分类模型，涵盖数据加载、模型定义、训练过程和结果评估。

1. 环境准备

首先，确保已安装PyTorch和必要的依赖库：

pip install torch torchvision matplotlib numpy

2. 数据准备与加载

2.1 数据集选择

本文使用经典的CIFAR-10数据集，包含10个类别的60000张32x32彩色图像（训练集50000张，测试集10000张）。

2.2 数据加载代码

import torch
import torchvision
import torchvision.transforms as transforms
# 定义数据预处理（归一化到[-1, 1]）
transform = transforms.Compose([
    transforms.ToTensor(),  # 将PIL图像或numpy数组转为Tensor，并缩放到[0, 1]
    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,  # 每次加载32张图像
    shuffle=True,  # 打乱数据顺序
    num_workers=2  # 使用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
)
# 类别名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

注释说明：

• transforms.Compose：将多个预处理操作组合在一起。
• ToTensor：将图像从PIL格式或numpy数组转为PyTorch Tensor，并自动缩放到[0, 1]。
• Normalize：使用均值和标准差对数据进行归一化，这里将像素值从[0, 1]缩放到[-1, 1]。
• DataLoader：封装数据集，提供批量加载、打乱和多线程加载功能。

3. 模型定义

3.1 卷积神经网络（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（RGB），输出通道6，卷积核大小5x5
        self.conv1 = nn.Conv2d(3, 6, 5)
        # 第二个卷积层：输入通道6，输出通道16，卷积核大小5x5
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 第一个全连接层：输入16*5*5（经过池化后的特征图大小），输出120
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        # 第二个全连接层：输入120，输出84
        self.fc2 = nn.Linear(120, 84)
        # 输出层：输入84，输出10（10个类别）
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        # 最大池化，kernel_size=2, stride=2
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # 最大池化，kernel_size=2, stride=2
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        # 展平特征图
        x = x.view(-1, 16 * 5 * 5)
        # 全连接层 + ReLU激活
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        # 输出层（不需要激活函数，因为后面会接CrossEntropyLoss）
        x = self.fc3(x)
        return x
# 实例化模型
net = Net()

注释说明：

• nn.Conv2d：定义卷积层，参数依次为输入通道数、输出通道数、卷积核大小。
• nn.Linear：定义全连接层，参数依次为输入特征数、输出特征数。
• F.max_pool2d：最大池化操作，参数依次为输入Tensor、池化核大小。
• F.relu：ReLU激活函数。
• x.view(-1, 16 * 5 * 5)：将特征图展平为一维向量，-1表示自动计算该维度大小。

4. 损失函数与优化器

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

注释说明：

• nn.CrossEntropyLoss：交叉熵损失，适用于多分类问题。
• optim.SGD：随机梯度下降优化器，net.parameters()表示优化模型的所有参数。

5. 模型训练

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)  # 将模型移动到GPU（如果可用）
for epoch in range(10):  # 训练10个epoch
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 获取输入和标签
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)  # 移动到GPU
        # 梯度清零
        optimizer.zero_grad()
        # 前向传播 + 反向传播 + 优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        # 打印统计信息
        running_loss += loss.item()
        if i % 1000 == 999:  # 每1000个batch打印一次
            print(f'Epoch {epoch + 1}, Batch {i + 1}, Loss: {running_loss / 1000:.3f}')
            running_loss = 0.0
print('Finished Training')

注释说明：

• torch.device：检查是否有可用的GPU。
• net.to(device)：将模型移动到GPU。
• optimizer.zero_grad()：清空上一步的梯度。
• loss.backward()：反向传播计算梯度。
• optimizer.step()：根据梯度更新参数。
• loss.item()：获取标量损失值。

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 = net(images)
        _, predicted = torch.max(outputs.data, 1)  # 获取概率最大的类别
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print(f'Accuracy on test set: {100 * correct / total:.2f}%')

注释说明：

• torch.no_grad()：上下文管理器，禁用梯度计算以节省内存和计算资源。
• torch.max(outputs.data, 1)：沿维度1（类别维度）取最大值，返回值和索引。
• predicted == labels：计算预测正确的样本数。

7. 完整代码与运行结果

7.1 完整代码

将上述代码片段整合为一个完整的Python脚本（见附录）。

7.2 运行结果示例

Epoch 1, Batch 1000, Loss: 2.189
Epoch 1, Batch 2000, Loss: 1.856
...
Epoch 10, Batch 1000, Loss: 0.456
Finished Training
Accuracy on test set: 62.34%

8. 总结与改进建议

8.1 总结

本文通过完整的代码和详细的注释，展示了如何使用PyTorch实现一个基础的图像分类模型。关键步骤包括：

1. 数据加载与预处理。
2. 定义CNN模型结构。
3. 选择损失函数和优化器。
4. 训练模型并监控损失。
5. 测试模型并评估准确率。

8.2 改进建议

1. 模型优化：尝试更深的网络结构（如ResNet），或使用预训练模型（如ResNet18）。
2. 数据增强：在数据预处理中加入随机裁剪、旋转等操作，提升模型泛化能力。
3. 超参数调优：调整学习率、批量大小等超参数，或使用学习率调度器。
4. 分布式训练：对于大规模数据集，可以使用多GPU或分布式训练加速。

附录：完整代码

# 完整代码（整合上述所有片段）
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
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=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
)
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.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 = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        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()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
# 定义损失函数和优化器
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
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 1000 == 999:
            print(f'Epoch {epoch + 1}, Batch {i + 1}, Loss: {running_loss / 1000:.3f}')
            running_loss = 0.0
print('Finished Training')
# 测试模型
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print(f'Accuracy on test set: {100 * correct / total:.2f}%')

通过本文的指导，读者可以快速上手PyTorch图像分类任务，并为后续更复杂的计算机视觉项目打下基础。