import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
===================== 1. 基础配置（极简版） =====================
设备：优先GPU，无则CPU（CPU也能快速运行）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
超参数：极致简化，减少耗时
BATCH_SIZE = 32 # 批次缩小，减少内存占用
EPOCHS = 3 # 仅训练3轮，几十秒完成
LR = 0.001 # 学习率适中
===================== 2. 数据加载（自动下载+极简预处理） =====================
仅做「转张量 + 归一化」，无多余操作
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)) # MNIST通用归一化参数
])
加载MNIST数据集（自动下载，约100MB，首次运行稍慢，后续秒加载）
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_data = datasets.MNIST('./data', train=False, download=True, transform=transform)
数据加载器：简化参数，无多线程（避免系统适配问题）
train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False)
===================== 3. 极简模型（1卷积+2全连接） =====================
class SimpleMNIST(nn.Module):
def init(self):
super().init()
# 卷积层：提取基础特征（28x28→14x14）
self.conv = nn.Sequential(
nn.Conv2d(1, 16, 3, padding=1), # 输入1通道，输出16通道
nn.ReLU(),
nn.MaxPool2d(2)
)
# 全连接层：分类（161414=3136 → 10类）
self.fc = nn.Sequential(
nn.Flatten(), # 展平卷积结果
nn.Linear(161414, 128), # 隐藏层
nn.ReLU(),
nn.Linear(128, 10) # 输出0-9
)
def forward(self, x):x = self.conv(x)x = self.fc(x)return x
初始化模型并移至设备
model = SimpleMNIST().to(device)
===================== 4. 训练+测试（极简流程） =====================
损失函数+优化器（基础配置，无多余操作）
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LR)
训练循环（极简版，无进度条，仅打印关键信息）
print("开始训练（仅3轮，很快完成）...")
for epoch in range(EPOCHS):
model.train() # 训练模式
for data, target in train_loader:
data, target = data.to(device), target.to(device)
# 前向+反向+更新（核心三步）
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
print(f"第{epoch+1}轮训练完成")
测试精度（极简版，仅计算整体准确率）
model.eval() # 验证模式
correct = 0
total = 0
with torch.no_grad(): # 禁用梯度，加速计算
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
_, pred = torch.max(output, 1) # 取预测结果
correct += (pred == target).sum().item()
total += target.size(0)
输出最终结果
acc = 100 * correct / total
print(f"\n训练完成！测试集准确率：{acc:.2f}%")
可选：单样本预测演示（验证模型可用）
print("\n单样本预测演示：")
test_img, test_label = test_data[0] # 取第一个测试样本
test_img = test_img.unsqueeze(0).to(device) # 增加批次维度
model.eval()
with torch.no_grad():
pred = model(test_img)
pred_label = torch.argmax(pred).item()
print(f"真实标签：{test_label}，模型预测：{pred_label}")
print("\n学号：2024310143004")