[Pytorch案例实践002]CIFAR-10图像分类

一、实践案例介绍

        这个项目的主要目标是使用 PyTorch 构建、训练和验证一个简单的卷积神经网络（CNN）模型，以在 CIFAR-10 数据集上进行图像分类任务。具体步骤如下：        

1. 设置设备：

   检测是否有可用的 GPU，并设置计算设备为 GPU 或 CPU。使用 GPU 可以加速模型训练过程。

2. 数据预处理：

   使用 `torchvision.transforms` 对 CIFAR-10 数据集进行预处理，将图像数据转换为 PyTorch 的张量，并进行归一化处理。

3. 加载数据集：

   使用 `torchvision.datasets.CIFAR10` 下载并加载 CIFAR-10 训练集和测试集。

   使用 `DataLoader` 将数据集包装成可迭代对象，便于批量处理和打乱数据。

4. 定义模型：

   构建一个简单的卷积神经网络 `SimpleCNN`，包含三个卷积层、最大池化层和全连接层。

   每个卷积层后接 ReLU 激活函数，以引入非线性特性。

5. 定义训练和验证过程：

   使用交叉熵损失函数（`nn.CrossEntropyLoss`）和 Adam 优化器（`optim.Adam`）。

   在训练阶段：

     逐批获取训练数据，进行前向传播、计算损失、反向传播并更新模型参数。

     使用 `tqdm` 库显示训练进度和当前的损失及准确率。

   在验证阶段：

     评估模型在测试集上的性能，计算测试损失和准确率。

     如果当前模型的准确率超过之前的最佳准确率，则保存该模型的参数。

6. 绘制结果：

   训练完成后，绘制并保存训练和验证过程中的损失曲线和准确率曲线，以直观展示模型的性能变化。

二、数据集介绍

        CIFAR-10 数据集是一个常用于计算机视觉领域的图像分类基准数据集，由加拿大高级研究院 (CIFAR) 提供。以下是 CIFAR-10 数据集的详细介绍：

数据集概述

• 数据集规模：包含 60000 张 32x32 彩色图像

  • 训练集：50000 张图像
  • 测试集：10000 张图像

• 类别：共 10 个类别，每个类别包含 6000 张图像

  • 飞机 (airplane)
  • 汽车 (automobile)
  • 鸟 (bird)
  • 猫 (cat)
  • 鹿 (deer)
  • 狗 (dog)
  • 青蛙 (frog)
  • 马 (horse)
  • 船 (ship)
  • 卡车 (truck)

数据集特点

• 图像尺寸：每张图像的尺寸为 32x32 像素。
• 图像类型：RGB 彩色图像。
• 标注：每张图像都被标注为 10 个类别中的一个。
• 平衡性：每个类别包含相同数量的图像，确保数据集的平衡性。

数据集用途

CIFAR-10 数据集广泛用于以下研究和任务：

• 图像分类：对图像进行分类，判断其所属类别。
• 计算机视觉算法评估：评估各种图像分类算法的性能。
• 深度学习模型训练：用于训练和测试深度学习模型，如卷积神经网络 (CNN)。

三、完整代码

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from tqdm import tqdm
import matplotlib.pyplot as plt

# 设置设备（使用GPU如果可用，否则使用CPU）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据预处理，转为Tensor并归一化
transform = transforms.Compose([
    transforms.ToTensor(),  # 转为Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # 归一化
])

# 加载CIFAR-10训练集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform)
trainloader = DataLoader(trainset, batch_size=8, shuffle=True, num_workers=2)

# 加载CIFAR-10测试集
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform)
testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=2)

# 构建简单的卷积神经网络
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),  # 卷积层1
            nn.ReLU(),  # 激活函数
            nn.MaxPool2d(kernel_size=2, stride=2),  # 池化层1
            nn.Conv2d(32, 64, kernel_size=3, padding=1),  # 卷积层2
            nn.ReLU(),  # 激活函数
            nn.MaxPool2d(kernel_size=2, stride=2),  # 池化层2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),  # 卷积层3
            nn.ReLU(),  # 激活函数
            nn.MaxPool2d(kernel_size=2, stride=2),  # 池化层3
            nn.Flatten(),  # 展平
            nn.Linear(128*4*4, 256),  # 全连接层1
            nn.ReLU(),  # 激活函数
            nn.Linear(256, 10)  # 全连接层2（输出层）
        )

    def forward(self, x):
        return self.net(x)

# 定义训练和验证的主函数
def main():
    # 实例化网络并移动到设备（GPU或CPU）
    net = SimpleCNN().to(device)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(), lr=0.001)

    # 训练和验证
    num_epochs = 10  # 训练的轮数
    best_acc = 0.0  # 初始化最佳准确率
    train_losses = []  # 用于保存训练损失
    test_losses = []  # 用于保存测试损失
    train_accuracies = []  # 用于保存训练准确率
    test_accuracies = []  # 用于保存测试准确率

    for epoch in range(num_epochs):
        net.train()  # 设置网络为训练模式
        running_loss = 0.0
        correct = 0
        total = 0

        # 使用tqdm显示训练进度条
        train_bar = tqdm(trainloader, desc=f'Training Epoch {epoch+1}/{num_epochs}')
        for inputs, labels in train_bar:
            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()
            _, predicted = outputs.max(1)  # 获取预测结果
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

            # 更新进度条信息
            train_bar.set_postfix(loss=running_loss/(len(train_bar)*trainloader.batch_size), acc=100.*correct/total)

        train_loss = running_loss / len(trainloader)
        train_acc = 100. * correct / total
        train_losses.append(train_loss)
        train_accuracies.append(train_acc)

        net.eval()  # 设置网络为评估模式
        test_loss = 0.0
        correct = 0
        total = 0
        with torch.no_grad():  # 禁用梯度计算
            # 使用tqdm显示验证进度条
            test_bar = tqdm(testloader, desc=f'Validating Epoch {epoch+1}/{num_epochs}')
            for inputs, labels in test_bar:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = net(inputs)
                loss = criterion(outputs, labels)
                test_loss += loss.item()
                _, predicted = outputs.max(1)
                total += labels.size(0)
                correct += predicted.eq(labels).sum().item()
                test_bar.set_postfix(loss=test_loss/(len(test_bar)*testloader.batch_size), acc=100.*correct/total)

        test_loss = test_loss / len(testloader)
        test_acc = 100. * correct / total
        test_losses.append(test_loss)
        test_accuracies.append(test_acc)

        # 保存最好的模型
        if test_acc > best_acc:
            best_acc = test_acc
            torch.save(net.state_dict(), 'best_model.pth')

    # 绘制损失和准确率曲线并保存
    plt.figure(figsize=(12, 5))

    # 绘制损失曲线
    plt.subplot(1, 2, 1)
    plt.plot(range(num_epochs), train_losses, label='Training Loss')
    plt.plot(range(num_epochs), test_losses, label='Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.title('Loss Curve')

    # 绘制准确率曲线
    plt.subplot(1, 2, 2)
    plt.plot(range(num_epochs), train_accuracies, label='Training Accuracy')
    plt.plot(range(num_epochs), test_accuracies, label='Validation Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.title('Accuracy Curve')

    # 保存图像
    plt.savefig('loss_accuracy_curves.png')
    plt.show()

if __name__ == "__main__":
    main()

四、总结

        本项目使用 PyTorch 构建、训练和验证一个简单的卷积神经网络，以实现 CIFAR-10 数据集上的图像分类。通过数据预处理、模型定义、训练和验证、性能评估等步骤，展示了如何在实际应用中进行深度学习模型的开发与优化。