手写数字识别

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

设备配置：优先使用GPU，无GPU则用CPU

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

1. 数据预处理

transform = transforms.Compose([
transforms.ToTensor(), # 转为张量并归一化到[0,1]
transforms.Normalize((0.1307,), (0.3081,)) # MNIST数据集的均值和标准差
])

加载MNIST数据集

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=1000, shuffle=False)

2. 定义卷积神经网络模型

class CNN(nn.Module):
def init(self):
super(CNN, self).init()
# 卷积层1: 1输入通道, 32输出通道, 3x3卷积核
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
# 卷积层2: 32输入通道, 64输出通道, 3x3卷积核
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
# 池化层: 2x2池化核
self.pool = nn.MaxPool2d(2, 2)
# 全连接层1: 6477 -> 128
self.fc1 = nn.Linear(64 * 7 * 7, 128)
# 全连接层2: 128 -> 10 (10个数字类别)
self.fc2 = nn.Linear(128, 10)
# Dropout层: 防止过拟合
self.dropout = nn.Dropout(0.25)

def forward(self, x):
    # 卷积1 -> 激活 -> 池化
    x = self.pool(F.relu(self.conv1(x)))
    # 卷积2 -> 激活 -> 池化
    x = self.pool(F.relu(self.conv2(x)))
    # 展平: [batch_size, 64, 7, 7] -> [batch_size, 64*7*7]
    x = x.view(-1, 64 * 7 * 7)
    # 全连接1 -> 激活 -> Dropout
    x = F.relu(self.fc1(x))
    x = self.dropout(x)
    # 全连接2 -> 输出
    x = self.fc2(x)
    return x

3. 初始化模型、损失函数、优化器

model = CNN().to(device)
criterion = nn.CrossEntropyLoss() # 交叉熵损失
optimizer = optim.Adam(model.parameters(), lr=0.001) # Adam优化器

4. 训练函数

def train(model, train_loader, criterion, optimizer, epoch):
model.train()
running_loss = 0.0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)

    # 前向传播
    output = model(data)
    loss = criterion(output, target)
    
    # 反向传播 + 优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    running_loss += loss.item()
    # 每100批次打印一次信息
    if batch_idx % 100 == 0:
        print(f'Epoch [{epoch+1}], Batch [{batch_idx}/{len(train_loader)}], Loss: {running_loss/100:.4f}')
        running_loss = 0.0

5. 测试函数

def test(model, test_loader):
model.eval()
test_loss = 0.0
correct = 0
with torch.no_grad(): # 测试阶段不计算梯度
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += criterion(output, target).item()
# 预测结果: 取概率最大的类别
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()

test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')
return accuracy

6. 主训练流程

if name == 'main':
# 输出学号
student_id = 3012
print(f"========== 学号: {student_id} ==========\n")

num_epochs = 5  # 训练轮数，可根据需要调整
for epoch in range(num_epochs):
    train(model, train_loader, criterion, optimizer, epoch)
    acc = test(model, test_loader)

print(f"训练完成！最终测试准确率: {acc:.2f}%")
print(f"学号: {student_id} | 手写数字识别任务完成")