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手写数字的高准确率分类能力。
模型结构可视化
该CNN采用两层卷积+池化提取特征,后接全连接层进行分类。卷积层使用ReLU激活函数引入非线性,Dropout层防止过拟合,最终输出10个类别的概率分布。