import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
# 1. 数据预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)) # MNIST单通道归一化
])
train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_data = datasets.MNIST(root='./data', train=False, transform=transform)
# 2. 模型定义
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.layers = nn.Sequential(
nn.Linear(28*28, 128), # 输入层
nn.ReLU(), # 激活函数
nn.Dropout(0.2), # 防过拟合
nn.Linear(128, 10) # 输出层（10分类）
)
def forward(self, x):
x = self.flatten(x)
return self.layers(x)
# 3. 训练配置
model = NeuralNetwork()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
batch_size = 64
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
# 4. 训练循环
for epoch in range(10):
for images, labels in train_loader:
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 5. 评估
test_loader = torch.utils.data.DataLoader(test_data, batch_size=256)
correct = 0
with torch.no_grad():
for images, labels in test_loader:
outputs = model(images)
_, predicted = torch.max(outputs, 1)
correct += (predicted == labels).sum().item()
print(f'准确率: {100 * correct / len(test_data):.2f}%')