PyTorch 2.0 CNN 图像分类实战：FashionMNIST 数据集 3 轮训练达 88% 准确率

PyTorch 2.0 实战：3轮训练实现FashionMNIST图像分类88%准确率

当我们需要快速验证一个图像分类模型的有效性时，FashionMNIST数据集和PyTorch的组合往往能提供绝佳的实验平台。这个经典的服装分类数据集不仅体积小巧适合快速迭代，还保留了真实图像数据的核心特征。本文将展示如何用PyTorch 2.0构建一个精简但高效的CNN模型，仅用3个训练周期就达到88%的测试准确率。

1. 环境准备与数据加载

让我们从最基础的环境配置开始。确保已安装PyTorch 2.0或更高版本，这是体验最新性能优化的关键。虽然本文示例可以在CPU上运行，但如果有CUDA兼容的GPU，训练速度将显著提升。

import torch from torch import nn import torchvision from torchvision import datasets, transforms from torch.utils.data import DataLoader # 检查PyTorch版本和设备 print(f"PyTorch版本: {torch.__version__}") device = "cuda" if torch.cuda.is_available() else "cpu"

FashionMNIST数据集包含10类服装的灰度图像，每张28x28像素。PyTorch内置了这个数据集，我们可以直接下载并使用：

# 定义数据转换 transform = transforms.Compose([ 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 ) # 创建数据加载器 BATCH_SIZE = 32 train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True) test_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False) # 查看类别名称 class_names = train_data.classes print(f"类别列表: {class_names}")

提示：将shuffle设为True可以防止模型学习到数据顺序带来的偏差，这对训练效果至关重要。

2. CNN模型架构设计

我们的目标是构建一个足够简单但有效的CNN结构。这个模型将包含两个卷积块，每个块由卷积层、ReLU激活和最大池化组成，最后接全连接层进行分类。

class FashionMNIST_CNN(nn.Module): def __init__(self, input_shape: int, hidden_units: int, output_shape: int): super().__init__() self.conv_block_1 = nn.Sequential( nn.Conv2d(in_channels=input_shape, out_channels=hidden_units, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2) ) self.conv_block_2 = nn.Sequential( nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2) ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(in_features=hidden_units*7*7, # 为什么是7*7？两次池化后28x28→14x14→7x7 out_features=output_shape) ) def forward(self, x): x = self.conv_block_1(x) x = self.conv_block_2(x) x = self.classifier(x) return x # 初始化模型 model = FashionMNIST_CNN(input_shape=1, # 灰度图单通道 hidden_units=10, output_shape=len(class_names)).to(device)

这个架构虽然简单，但包含了CNN的核心要素：

• 局部感受野：3x3卷积核捕捉局部特征
• 非线性激活：ReLU引入非线性表达能力
• 下采样：最大池化减少计算量同时保持特征不变性
• 全连接分类：最后将特征映射到10个类别

3. 训练流程实现

训练循环是深度学习的核心引擎。我们将定义损失函数、优化器，并实现完整的训练和测试步骤。

# 设置损失函数和优化器 loss_fn = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(params=model.parameters(), lr=0.1) # 定义准确率计算函数 def accuracy_fn(y_true, y_pred): correct = torch.eq(y_true, y_pred).sum().item() acc = (correct / len(y_pred)) * 100 return acc # 训练步骤 def train_step(model: nn.Module, data_loader: DataLoader, loss_fn: nn.Module, optimizer: torch.optim.Optimizer, accuracy_fn, device: torch.device = device): train_loss, train_acc = 0, 0 model.train() for batch, (X, y) in enumerate(data_loader): X, y = X.to(device), y.to(device) # 前向传播 y_pred = model(X) # 计算损失 loss = loss_fn(y_pred, y) train_loss += loss train_acc += accuracy_fn(y_true=y, y_pred=y_pred.argmax(dim=1)) # 优化器清零梯度 optimizer.zero_grad() # 反向传播 loss.backward() # 参数更新 optimizer.step() # 计算平均损失和准确率 train_loss /= len(data_loader) train_acc /= len(data_loader) print(f"训练损失: {train_loss:.5f} | 训练准确率: {train_acc:.2f}%") # 测试步骤 def test_step(model: nn.Module, data_loader: DataLoader, loss_fn: nn.Module, accuracy_fn, device: torch.device = device): test_loss, test_acc = 0, 0 model.eval() with torch.inference_mode(): for X, y in data_loader: X, y = X.to(device), y.to(device) # 前向传播 test_pred = model(X) # 计算损失和准确率 test_loss += loss_fn(test_pred, y) test_acc += accuracy_fn(y_true=y, y_pred=test_pred.argmax(dim=1)) # 计算平均损失和准确率 test_loss /= len(data_loader) test_acc /= len(data_loader) print(f"测试损失: {test_loss:.5f} | 测试准确率: {test_acc:.2f}%\n")

4. 模型训练与性能评估

现在让我们启动训练流程，仅用3个epoch观察模型表现：

# 设置随机种子保证可重复性 torch.manual_seed(42) # 训练3个epoch epochs = 3 for epoch in range(epochs): print(f"Epoch {epoch+1}\n--------") train_step(model=model, data_loader=train_loader, loss_fn=loss_fn, optimizer=optimizer, accuracy_fn=accuracy_fn, device=device) test_step(model=model, data_loader=test_loader, loss_fn=loss_fn, accuracy_fn=accuracy_fn, device=device)

典型输出结果可能如下：

Epoch 1 -------- 训练损失: 0.58912 | 训练准确率: 78.67% 测试损失: 0.42601 | 测试准确率: 84.62% Epoch 2 -------- 训练损失: 0.38710 | 训练准确率: 86.10% 测试损失: 0.37610 | 测试准确率: 86.45% Epoch 3 -------- 训练损失: 0.34572 | 训练准确率: 87.47% 测试损失: 0.35821 | 测试准确率: 87.72%

从结果可以看到，仅用3个训练周期，模型在测试集上的准确率就达到了接近88%的水平。这种快速收敛的特性使得这个配置非常适合作为更复杂模型的基准参考。

5. 关键参数与调优建议

虽然我们的基础模型已经表现不错，但理解关键参数的影响对进一步优化至关重要：

# 示例：增强版模型架构 class Enhanced_CNN(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(1, 32, 3, padding=1) self.conv2 = nn.Conv2d(32, 64, 3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.dropout = nn.Dropout(0.25) self.fc1 = nn.Linear(64*7*7, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = self.dropout(x) x = x.view(-1, 64*7*7) x = F.relu(self.fc1(x)) x = self.dropout(x) x = self.fc2(x) return x

这个增强版引入了Dropout正则化和更多的卷积通道，通常能获得更好的表现，但训练时间也会相应增加。