PyTorch实战：从零开始构建CIFAR-10图像分类模型 (附详细代码与图解)

PyTorch实战：从零开始构建CIFAR-10图像分类模型 (附详细代码与图解)

大家好！今天，我们将一起踏上一段激动人心的深度学习之旅：使用强大的PyTorch框架，从零开始构建一个卷积神经网络（CNN），来解决经典的CIFAR-10图像分类问题。

无论你是深度学习的新手，还是希望巩固PyTorch基础知识的开发者，本文都将为你提供一个清晰、详尽的实战指南。

本文目标

读完本文，你将学会：

1. 加载和预处理 CIFAR-10 数据集。
2. 使用PyTorch设计并构建一个自定义的卷积神经网络。
3. 定义损失函数和优化器。
4. 训练你的模型，并实时监控训练过程。
5. 在测试集上评估模型的性能，并分析结果。

整体流程概览

在深入代码之前，我们先用一张图来梳理整个项目的流程，让你有一个宏观的认识。

评估

训练

模型

数据

4. 评估与分析

4.1 在测试集上评估

4.2 分析各类别准确率

3. 训练模型

3.1 定义损失函数

3.2 定义优化器

3.3 编写训练循环

2. 构建模型

2.1 设计CNN网络结构

2.2 用PyTorch代码实现

1. 准备数据

1.1 加载CIFAR-10数据集

1.2 数据预处理与增强

1.3 创建DataLoader

第一步：准备工作 - 环境与数据

1.1 导入必要的库

首先，确保你已经安装了torch和torchvision。然后，我们导入所有需要的库。

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
import matplotlib.pyplot as plt
import numpy as np

1.2 认识CIFAR-10数据集

CIFAR-10是一个包含了10个类别彩色图像的数据集。

• 图像大小: 32x32像素
• 通道数: 3 (RGB)
• 类别: 10类，包括 ‘plane’, ‘car’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’
• 数量: 60000张图片，其中50000张用于训练，10000张用于测试。

1.3 数据加载与预处理

torchvision库让数据加载变得异常简单。在加载的同时，我们需要对数据进行预处理（transforms）：

1. ToTensor(): 将PIL Image或Numpy ndarray格式的图片转换为torch.Tensor，并将像素值从[0, 255]缩放到[0, 1]。
2. Normalize(mean, std): 对张量进行标准化。公式为 output = (input - mean) / std。标准化有助于模型更快、更稳定地收敛。这里我们使用mean=(0.5, 0.5, 0.5)和std=(0.5, 0.5, 0.5)，将数据从[0, 1]范围转换到[-1, 1]范围。

# 定义数据预处理
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 定义批量大小
batch_size = 64

# 下载并加载训练集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          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=batch_size,
                                         shuffle=False, num_workers=2)

# 10个类别的名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

1.4 可视化部分数据

让我们看几张训练图片，直观感受一下数据集。

def imshow(img):
    """显示图像的函数"""
    img = img / 2 + 0.5     # 反标准化，从[-1,1] -> [0,1]
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 获取一些随机的训练图像
dataiter = iter(trainloader)
images, labels = next(dataiter)

# 显示图像
imshow(torchvision.utils.make_grid(images[:8])) # 显示前8张图
# 打印标签
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(8)))

运行后，你会看到类似下面的输出和图像：

ship  horse truck plane frog  cat   deer  bird 

第二步：核心环节 - 构建卷积神经网络(CNN)

我们将构建一个简单的CNN模型，其结构如下：

Input -> Conv1 -> ReLU -> Pool1 -> Conv2 -> ReLU -> Pool2 -> Flatten -> FC1 -> ReLU -> FC2 -> ReLU -> FC3(Output)

下面是这个结构的图解：

graph TD
    subgraph CNN Architecture
        A[Input (3x32x32)] --> B(Conv1: 3x32x32 -> 6x28x28);
        B --> C(ReLU);
        C --> D(MaxPool1: 6x28x28 -> 6x14x14);
        
        D --> E(Conv2: 6x14x14 -> 16x10x10);
        E --> F(ReLU);
        F --> G(MaxPool2: 16x10x10 -> 16x5x5);
        
        G --> H(Flatten: 16x5x5 -> 400);
        
        H --> I(FC1: 400 -> 120);
        I --> J(ReLU);
        J --> K(FC2: 120 -> 84);
        K --> L(ReLU);
        L --> M(FC3: 84 -> 10 (Output));
    end

现在，我们用PyTorch代码来实现这个网络。

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        # 卷积层
        # 输入通道=3, 输出通道=6, 卷积核大小=5x5
        self.conv1 = nn.Conv2d(3, 6, 5)
        # 池化层, 窗口大小=2x2, 步长=2
        self.pool = nn.MaxPool2d(2, 2)
        # 输入通道=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):
        # x的shape: [batch_size, 3, 32, 32]
        
        # 卷积 -> 激活 -> 池化
        x = self.pool(F.relu(self.conv1(x))) # -> [batch_size, 6, 14, 14]
        x = self.pool(F.relu(self.conv2(x))) # -> [batch_size, 16, 5, 5]
        
        # 展平操作
        x = torch.flatten(x, 1) # -> [batch_size, 16 * 5 * 5] = [batch_size, 400]
        
        # 全连接层
        x = F.relu(self.fc1(x)) # -> [batch_size, 120]
        x = F.relu(self.fc2(x)) # -> [batch_size, 84]
        x = self.fc3(x)         # -> [batch_size, 10]
        return x

# 实例化模型并移到GPU(如果可用)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

net = Net().to(device)
print(net) # 打印网络结构

第三步：定义损失函数与优化器

• 损失函数 (Loss Function): 对于多分类问题，我们通常使用交叉熵损失（nn.CrossEntropyLoss）。这个函数内部已经包含了Softmax操作，所以我们的网络最后一层不需要加Softmax。
• 优化器 (Optimizer): 我们选择经典的随机梯度下降（SGD），并加上动量（momentum）来加速收敛。

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

第四步：开始训练！

这是最核心的部分。我们将遍历我们的数据集多次（称为epochs），在每个epoch中，我们分批次（batch）地将数据送入模型进行训练。

训练的步骤如下：

1. 从trainloader中获取一个批次的数据和标签。
2. 将数据和标签移到GPU。
3. 清零梯度 (optimizer.zero_grad())。
4. 前向传播 (outputs = net(inputs))。
5. 计算损失 (loss = criterion(outputs, labels))。
6. 反向传播 (loss.backward())。
7. 更新权重 (optimizer.step())。

print("Starting Training...")

num_epochs = 10  # 训练10个周期

for epoch in range(num_epochs):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 1. 获取输入数据
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        # 2. 梯度清零
        optimizer.zero_grad()

        # 3. 前向传播 + 反向传播 + 优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # 4. 打印统计信息
        running_loss += loss.item()
        if i % 200 == 199:    # 每200个mini-batches打印一次
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 200:.3f}')
            running_loss = 0.0

print('Finished Training')

# 保存模型
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

第五步：模型评估

训练完成了，但模型表现如何？我们需要在从未见过的测试集上进行评估。

5.1 整体准确率

我们将遍历整个testloader，计算模型的总体预测准确率。

correct = 0
total = 0
# 由于我们不是在训练，所以我们不需要计算梯度
net.eval() # 将模型设置为评估模式
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        
        # 前向传播
        outputs = net(images)
        
        # 获取预测结果
        # torch.max返回(最大值, 最大值索引)
        _, predicted = torch.max(outputs.data, 1)
        
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f'Accuracy of the network on the 10000 test images: {accuracy:.2f} %')

对于这个简单的网络，你可能会得到50%-65%左右的准确率。虽然不高，但作为一个起点已经很棒了！

5.2 各类别准确率分析

整体准确率可能会掩盖一些问题，比如模型可能对某些类别识别得很好，对另一些则很差。让我们来分别看看每个类别的表现。

# 准备统计每个类别的预测情况
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}

net.eval()
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = net(images)
        _, predictions = torch.max(outputs, 1)
        
        # 收集每个类别的正确预测
        for label, prediction in zip(labels, predictions):
            if label == prediction:
                correct_pred[classes[label]] += 1
            total_pred[classes[label]] += 1

# 打印每个类别的准确率
for classname, correct_count in correct_pred.items():
    accuracy = 100 * float(correct_count) / total_pred[classname]
    print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')

5.3 结果可视化

用一个条形图来展示各类别准确率会更直观。

class_accuracies = [100 * float(correct_pred[c]) / total_pred[c] for c in classes]

plt.figure(figsize=(12, 6))
plt.bar(classes, class_accuracies)
plt.xlabel('Class')
plt.ylabel('Accuracy (%)')
plt.title('Accuracy for each class on CIFAR-10')
plt.ylim([0, 100])
for i, acc in enumerate(class_accuracies):
    plt.text(i, acc + 1, f'{acc:.1f}', ha='center')
plt.show()

这张图清晰地告诉我们，模型在识别car和ship等物体上表现较好，但在识别cat和bird等动物上可能稍差。这为我们下一步改进模型提供了方向。

这是将整篇文章中的Python代码整合在一起的完整脚本。

您可以将此代码复制到一个.py文件中（例如 train_cifar10.py），然后直接运行。代码中包含了详细的注释，与博客文章的步骤一一对应。

# -----------------------------------------------------------------
# PyTorch实战：从零开始构建CIFAR-10图像分类模型
# 完整可执行脚本
# -----------------------------------------------------------------

# 1. 准备工作 - 导入必要的库
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
import matplotlib.pyplot as plt
import numpy as np

# -----------------------------------------------------------------
# 2. 数据加载与预处理
# -----------------------------------------------------------------

def load_data():
    """加载并预处理CIFAR-10数据集"""
    print("正在加载和预处理数据...")
    
    # 定义数据预处理
    # 1. 将PIL Image或Numpy ndarray格式的图片转换为torch.Tensor
    # 2. 将像素值从[0, 255]缩放到[0, 1]，然后标准化到[-1, 1]
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    # 定义批量大小
    batch_size = 64

    # 下载并加载训练集
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                              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=batch_size,
                                             shuffle=False, num_workers=2)

    # 10个类别的名称
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 
               'dog', 'frog', 'horse', 'ship', 'truck')
    
    print("数据加载完毕。")
    return trainloader, testloader, classes

# -----------------------------------------------------------------
# (可选) 可视化部分数据
# -----------------------------------------------------------------
def imshow(img):
    """显示图像的函数"""
    img = img / 2 + 0.5     # 反标准化, 从[-1,1] -> [0,1]
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# -----------------------------------------------------------------
# 3. 构建卷积神经网络(CNN)
# -----------------------------------------------------------------
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        # 卷积层
        self.conv1 = nn.Conv2d(3, 6, 5)  # 输入:3通道, 输出:6通道, 卷积核:5x5
        self.pool = nn.MaxPool2d(2, 2)   # 2x2池化
        self.conv2 = nn.Conv2d(6, 16, 5) # 输入:6通道, 输出:16通道, 卷积核:5x5
        
        # 全连接层
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10) # 输出10个类别

    def forward(self, x):
        # x shape: [batch_size, 3, 32, 32]
        x = self.pool(F.relu(self.conv1(x))) # -> [batch_size, 6, 14, 14]
        x = self.pool(F.relu(self.conv2(x))) # -> [batch_size, 16, 5, 5]
        x = torch.flatten(x, 1)             # 展平 -> [batch_size, 16*5*5=400]
        x = F.relu(self.fc1(x))             # -> [batch_size, 120]
        x = F.relu(self.fc2(x))             # -> [batch_size, 84]
        x = self.fc3(x)                     # -> [batch_size, 10]
        return x

# -----------------------------------------------------------------
# 4. 训练模型
# -----------------------------------------------------------------
def train_model(net, trainloader, device, num_epochs=10):
    """训练模型"""
    print("开始训练...")
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
    
    net.train() # 将模型设置为训练模式
    for epoch in range(num_epochs):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            # 获取输入数据
            inputs, labels = data[0].to(device), data[1].to(device)

            # 梯度清零
            optimizer.zero_grad()

            # 前向传播 -> 反向传播 -> 优化
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            # 打印统计信息
            running_loss += loss.item()
            if i % 200 == 199:    # 每200个mini-batches打印一次
                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 200:.3f}')
                running_loss = 0.0

    print('训练完成。')
    # 保存模型
    PATH = './cifar_net.pth'
    torch.save(net.state_dict(), PATH)
    print(f'模型已保存到 {PATH}')


# -----------------------------------------------------------------
# 5. 评估模型
# -----------------------------------------------------------------
def evaluate_model(net, testloader, classes, device):
    """在测试集上评估模型"""
    print("开始在测试集上评估模型...")
    net.eval() # 将模型设置为评估模式

    # 5.1 评估整体准确率
    correct = 0
    total = 0
    with torch.no_grad(): # 在此模式下，不计算梯度
        for data in testloader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'模型在10000张测试图像上的准确率: {accuracy:.2f} %')

    # 5.2 评估每个类别的准确率
    correct_pred = {classname: 0 for classname in classes}
    total_pred = {classname: 0 for classname in classes}

    with torch.no_grad():
        for data in testloader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = net(images)
            _, predictions = torch.max(outputs, 1)
            for label, prediction in zip(labels, predictions):
                if label == prediction:
                    correct_pred[classes[label]] += 1
                total_pred[classes[label]] += 1

    # 打印并可视化每个类别的准确率
    class_accuracies = []
    print("\n各类别准确率:")
    for classname, correct_count in correct_pred.items():
        accuracy = 100 * float(correct_count) / total_pred[classname]
        class_accuracies.append(accuracy)
        print(f'类别: {classname:5s} 的准确率是 {accuracy:.1f} %')
    
    # 5.3 可视化结果
    plt.figure(figsize=(12, 6))
    plt.bar(classes, class_accuracies, color='skyblue')
    plt.xlabel('类别')
    plt.ylabel('准确率 (%)')
    plt.title('CIFAR-10各类别预测准确率')
    plt.ylim([0, 100])
    for i, acc in enumerate(class_accuracies):
        plt.text(i, acc + 1, f'{acc:.1f}', ha='center', color='black')
    plt.show()

# -----------------------------------------------------------------
# 主函数
# -----------------------------------------------------------------
if __name__ == '__main__':
    # 设置设备 (GPU or CPU)
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")

    # 1. 加载数据
    trainloader, testloader, classes = load_data()

    # (可选) 显示一些图像
    # print("显示部分训练图像...")
    # dataiter = iter(trainloader)
    # images, labels = next(dataiter)
    # imshow(torchvision.utils.make_grid(images[:8]))
    # print(' '.join(f'{classes[labels[j]]:5s}' for j in range(8)))

    # 2. 实例化模型并移到设备
    net = Net().to(device)
    print("\n模型结构:")
    print(net)

    # 3. 训练模型
    train_model(net, trainloader, device, num_epochs=10)

    # 4. 评估模型
    # 注意：如果跳过训练，需要先加载已保存的模型
    # net = Net().to(device)
    # net.load_state_dict(torch.load('./cifar_net.pth'))
    evaluate_model(net, testloader, classes, device)

总结与展望

恭喜你！你已经成功地使用PyTorch构建、训练并评估了一个完整的图像分类模型。我们回顾一下整个流程：

1. 数据准备：使用torchvision加载并预处理CIFAR-10。
2. 模型构建：设计了一个包含卷积、池化和全连接层的CNN。
3. 训练：定义了损失函数和优化器，并编写了标准的训练循环。
4. 评估：计算了整体准确率和各类别准确率，并进行了可视化分析。

如何更进一步？

这个模型只是一个起点，想要获得更高的准确率，你可以尝试：

• 数据增强：在transforms中加入随机翻转、裁剪、颜色抖动等操作。
• 更深、更复杂的网络：尝试VGG、ResNet等更强大的网络结构。
• 使用预训练模型：在ImageNet上预训练的模型进行迁移学习，会大大提升性能。
• 调整超参数：尝试不同的学习率、优化器（如Adam）、批次大小和训练周期。

希望这篇详细的教程能帮助你更好地理解和使用PyTorch。动手实践是最好的学习方式，快去敲代码，亲自体验一下吧！如果你有任何问题，欢迎在评论区留言交流。

---