[Pytorch案例实践003]Fashion-MNIST 图像分类

一、项目概述

        本项目旨在使用 PyTorch 框架构建和训练一个卷积神经网络（CNN），用于对 Fashion-MNIST 数据集进行分类。Fashion-MNIST 数据集是一个流行的图像分类数据集，包含 10 类不同的服装物品，每个类别有 7000 张灰度图像。每张图像的尺寸为 28x28 像素。

1. 数据预处理：

   • 加载并预处理 Fashion-MNIST 数据集，包括将图像转换为 Tensor 并进行归一化处理。

2. 模型构建：

   • 使用 nn.Sequential 构建一个简单的卷积神经网络，包括卷积层、ReLU 激活函数、最大池化层、展平层和全连接层。

3. 训练和验证：

   • 实现模型的训练和验证过程，包括计算损失和准确率。
   • 使用 tqdm 库显示训练和验证过程中的进度条。
   • 保存训练和验证集的损失和准确率曲线。

4. 模型保存：

   • 保存训练过程中表现最佳的模型权重。

5. 结果可视化：

   • 绘制并保存训练和验证过程中的损失和准确率曲线图。

二、数据集介绍

概述

Fashion-MNIST 数据集是一个用于图像分类任务的替代数据集，由 Zalando 提供。与经典的 MNIST 数据集不同，Fashion-MNIST 包含了更具挑战性的服装图像，旨在作为计算机视觉和机器学习算法的基准测试。

数据集组成

• 类别：10 类不同的服装物品，包括 T 恤、裤子、外套、连衣裙等。
• 样本数量：共 70,000 张灰度图像，其中 60,000 张用于训练，10,000 张用于测试。
• 图像尺寸：每张图像为 28x28 像素。
• 图像格式：灰度图像（单通道）。

类别标签

每张图像对应一个类别标签，共有以下 10 个类别：

1. T 恤/上衣
2. 裤子
3. 套头衫
4. 连衣裙
5. 外套
6. 凉鞋
7. 衬衫
8. 运动鞋
9. 包
10. 短靴

示例图像

数据集特点

1. 高相似度：Fashion-MNIST 中的许多图像在视觉上具有相似性，这使得分类任务比 MNIST 更具挑战性。
2. 实际应用：相比 MNIST 的手写数字，Fashion-MNIST 的图像更加贴近实际生活场景，更能测试模型的实用性。
3. 标准基准：Fashion-MNIST 被广泛用于测试和比较不同的机器学习算法，尤其是深度学习模型。

应用场景

Fashion-MNIST 数据集常用于以下场景：

• 模型验证和评估：作为标准数据集，用于验证和评估新模型的性能。
• 图像分类任务：测试图像分类算法的有效性。
• 深度学习研究：测试和比较不同的神经网络结构和优化算法。

三、完整代码

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,))  # 归一化
])

# 加载Fashion-MNIST训练集
trainset = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=100, shuffle=True, num_workers=2)

# 加载Fashion-MNIST测试集
testset = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)
testloader = DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)

# 构建简单的卷积神经网络
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 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*3*3, 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 进行 Fashion-MNIST 分类任务的完整流程，从数据预处理、模型构建、训练和验证到结果可视化和模型保存，展示了卷积神经网络在图像分类任务中的应用。

• 项目亮点：

  • 使用 nn.Sequential 简化了模型定义过程。
  • 通过 tqdm 显示训练和验证进度，提升了用户体验。
  • 保存最佳模型权重，确保最终模型性能最优。
  • 绘制损失和准确率曲线，直观展示模型训练效果。

• 未来改进：

  • 可以尝试更复杂的网络结构，如增加卷积层和全连接层，以进一步提高模型性能。
  • 可以应用数据增强技术，提升模型的泛化能力。
  • 可以尝试不同的优化器和学习率调度策略，进一步优化训练过程。

        通过本项目的实践，你不仅掌握了 Fashion-MNIST 数据集的基本使用方法，还了解了如何构建、训练和评估卷积神经网络，具备了进行图像分类任务的基础能力。