数字识别

-- coding: utf-8 --

"""
Created on Thu Dec 25 23:12:35 2025

@author: cxy-fairytale
"""

import torch

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

1. 数据加载与预处理

transform = transforms.Compose([
transforms.ToTensor(), # 转为张量
transforms.Normalize((0.5,), (0.5,)) # 归一化到 [-1, 1]
])

加载 MNIST 数据集

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

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

2. 定义 CNN 模型

class SimpleCNN(nn.Module):
def init(self):
super(SimpleCNN, self).init()
# 定义卷积层
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) # 输入1通道，输出32通道
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) # 输入32通道，输出64通道
# 定义全连接层
self.fc1 = nn.Linear(64 * 7 * 7, 128) # 展平后输入到全连接层
self.fc2 = nn.Linear(128, 10) # 10 个类别

def forward(self, x):
    x = F.relu(self.conv1(x))  # 第一层卷积 + ReLU
    x = F.max_pool2d(x, 2)     # 最大池化
    x = F.relu(self.conv2(x))  # 第二层卷积 + ReLU
    x = F.max_pool2d(x, 2)     # 最大池化
    x = x.view(-1, 64 * 7 * 7) # 展平
    x = F.relu(self.fc1(x))    # 全连接层 + ReLU
    x = self.fc2(x)            # 最后一层输出
    return x

创建模型实例

model = SimpleCNN()

3. 定义损失函数与优化器

criterion = nn.CrossEntropyLoss() # 多分类交叉熵损失
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

4. 模型训练

num_epochs = 5
model.train() # 设置模型为训练模式

for epoch in range(num_epochs):
total_loss = 0
for images, labels in train_loader:
outputs = model(images) # 前向传播
loss = criterion(outputs, labels) # 计算损失

    optimizer.zero_grad()  # 清空梯度
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数

    total_loss += loss.item()

print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss / len(train_loader):.4f}")

5. 模型测试

model.eval() # 设置模型为评估模式
correct = 0
total = 0

with torch.no_grad(): # 关闭梯度计算
for images, labels in test_loader:
outputs = model(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy:.2f}%")
print('3028')