1. 图像分类基础

任务目标

构建并训练一个简单的CNN模型,在MNIST手写数字数据集上完成图像分类任务,最终测试集准确率目标不低于99%。

环境与依赖

需要Python 3.8+,并安装以下依赖:

pip install torch torchvision matplotlib

完整实现代码

以下代码完整实现了数据加载、CNN模型定义、训练与评估全流程:

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 1. 数据准备
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)

# 2. 定义CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 64 * 7 * 7)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

model = SimpleCNN()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 3. 训练配置
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 4. 训练循环
def train(epochs=5):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f'Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(train_loader):.4f}')

# 5. 评估函数
def evaluate():
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = 100 * correct / total
    print(f'测试集准确率: {accuracy:.2f}%')
    return accuracy

# 6. 执行训练与评估
if __name__ == '__main__':
    train(epochs=5)
    evaluate()

运行方式

将上述代码保存为 mnist_cnn.py,在终端执行:

python mnist_cnn.py

首次运行会自动下载MNIST数据集到 ./data 目录。训练过程约需2-5分钟(CPU)或更短(GPU)。

验证方法与通过标准

运行完成后观察终端输出:

  • 训练损失应随epoch增加而下降,最终收敛至0.05以下;
  • 测试集准确率应达到99%以上,若低于此值可增加训练轮次至10轮或调整学习率。

满足上述条件即视为模型训练成功,具备了对MNIST手写数字的高准确率分类能力。

模型结构可视化

输入 28×28×1 Conv1 3×3, 32 → 28×28×32 Pool1 → 14×14×32 Conv2 3×3, 64 → 14×14×64 Pool2 → 7×7×64 FC 128→10

该CNN采用两层卷积+池化提取特征,后接全连接层进行分类。卷积层使用ReLU激活函数引入非线性,Dropout层防止过拟合,最终输出10个类别的概率分布。