当前位置：首页 > news >正文

基于Python的人工智能应用案例系列（14）：Fashion MNIST图像分类CNN

news 2025/7/4 15:27:27

在这一篇文章中，我们将使用PyTorch来实现卷积神经网络（CNN），对Fashion MNIST数据集进行图像分类任务。Fashion MNIST数据集是MNIST的升级版，包含各种服装、鞋类和配饰的灰度图像，非常适合作为深度学习的入门数据集。

数据集概述

Fashion MNIST数据集由60,000张训练图像和10,000张测试图像组成，每张图像的尺寸为28x28像素，且为灰度图像。数据集共包含10个类别，分别为：

T-shirt/top（T恤/上衣）
Trouser（裤子）
Pullover（套衫）
Dress（连衣裙）
Coat（外套）
Sandal（凉鞋）
Shirt（衬衫）
Sneaker（运动鞋）
Bag（包）
Ankle boot（短靴）

接下来，我们将从数据的加载开始，一步步构建并训练CNN模型。

1. ETL（提取、转换、加载）

首先，我们通过torchvision中的datasets模块加载Fashion MNIST数据集，并使用DataLoader进行批量处理。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import make_grid
import matplotlib.pyplot as plt# 数据转换
transform = transforms.ToTensor()# 加载训练和测试数据集
train_data = datasets.FashionMNIST(root='../Data', train=True, download=True, transform=transform)
test_data = datasets.FashionMNIST(root='../Data', train=False, download=True, transform=transform)# 定义类别名称
class_names = ['T-shirt', 'Trouser', 'Sweater', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Boot']# 使用DataLoader加载数据
train_loader = DataLoader(train_data, batch_size=10, shuffle=True)
test_loader = DataLoader(test_data, batch_size=10, shuffle=False)

2. EDA（探索性数据分析）

在训练模型之前，我们可以通过批量可视化数据来对数据集有一个直观的认识。

for images, labels in train_loader: breakprint('标签: ', labels.numpy())
print('类别: ', *[class_names[i] for i in labels])# 显示图像
im = make_grid(images, nrow=10)
plt.figure(figsize=(12,4))
plt.imshow(np.transpose(im.numpy(), (1, 2, 0)))

3. 模型训练

接下来我们定义一个卷积神经网络模型，该模型包含两个卷积层、两个池化层和两个全连接层。

class ConvolutionalNetwork(nn.Module):def __init__(self):super().__init__()self.conv1 = nn.Conv2d(1, 6, 3, 1)  # 输入通道数为1（灰度图），输出通道数为6，卷积核大小为3x3self.conv2 = nn.Conv2d(6, 16, 3, 1)  # 第二个卷积层，输出通道数为16self.fc1 = nn.Linear(5*5*16, 100)  # 全连接层，输入尺寸根据卷积结果计算self.fc2 = nn.Linear(100, 10)  # 输出层，10个类别def forward(self, X):X = F.relu(self.conv1(X))X = F.max_pool2d(X, 2, 2)X = F.relu(self.conv2(X))X = F.max_pool2d(X, 2, 2)X = X.view(-1, 5*5*16)X = F.relu(self.fc1(X))X = self.fc2(X)return Xtorch.manual_seed(101)
model = ConvolutionalNetwork()

模型参数统计

我们可以通过函数统计模型的可训练参数数量。

def count_parameters(model):params = [p.numel() for p in model.parameters() if p.requires_grad]for item in params:print(f'{item:>6}')print(f'总参数数量: {sum(params):>6}')count_parameters(model)

定义损失函数与优化器

我们使用交叉熵损失函数（CrossEntropyLoss）和Adam优化器。

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

模型训练

接下来我们开始训练模型，设定训练5个epoch。

epochs = 5for i in range(epochs):for X_train, y_train in train_loader:# 应用模型y_pred = model(X_train)loss = criterion(y_pred, y_train)# 更新参数optimizer.zero_grad()loss.backward()optimizer.step()print(f'{i+1} / {epochs} 轮训练完成')

4. 模型测试

训练完成后，我们对测试数据进行推理并计算准确率。

model.eval()with torch.no_grad():correct = 0for X_test, y_test in test_loader:y_val = model(X_test)predicted = torch.max(y_val,1)[1]correct += (predicted == y_test).sum()print(f'测试准确率: {correct.item()}/{len(test_data)} = {correct.item()*100/(len(test_data)):.3f}%')