手写数字识别

import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")
print("3002") # 新增：输出3002

BATCH_SIZE = 64
EPOCHS = 5
LEARNING_RATE = 0.001
IMAGE_SIZE = 28

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=BATCH_SIZE,
shuffle=True
)

test_loader = DataLoader(
test_dataset,
batch_size=BATCH_SIZE,
shuffle=False
)

def show_sample():
sample_image, sample_label = train_dataset[0]
sample_image = sample_image.squeeze().numpy()

plt.figure(figsize=(4,4))
plt.imshow(sample_image, cmap='gray')
plt.title(f"Label: {sample_label}")
plt.axis('off')
plt.show()

show_sample()

class MNIST_CNN(nn.Module):
def init(self):
super(MNIST_CNN, self).init()
self.conv_layers = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2, 2),

nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2, 2),
)

self.fc_layers = nn.Sequential(
nn.Linear(64 * 7 * 7, 128),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(128, 10)
)

def forward(self, x):
x = self.conv_layers(x)
x = x.view(x.size(0), -1)
x = self.fc_layers(x)
return x

model = MNIST_CNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

def train_model():
model.train()
total_loss = 0
correct = 0
total = 0

for batch_idx, (images, labels) in enumerate(train_loader):
images, labels = images.to(device), labels.to(device)

outputs = model(images)
loss = criterion(outputs, labels)

optimizer.zero_grad()
loss.backward()
optimizer.step()

total_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()

if (batch_idx + 1) % 100 == 0:
print(f'批次 [{batch_idx+1}/{len(train_loader)}], '
f'损失: {loss.item():.4f}, '
f'准确率: {(100 * correct / total):.2f}%')

avg_loss = total_loss / len(train_loader)
avg_acc = 100 * correct / total
return avg_loss, avg_acc

def test_model():
model.eval()
total_loss = 0
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)
loss = criterion(outputs, labels)

total_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()

avg_loss = total_loss / len(test_loader)
avg_acc = 100 * correct / total
return avg_loss, avg_acc

best_test_acc = 0.0
print("\n开始训练...")
for epoch in range(EPOCHS):
print(f"\n第 {epoch+1}/{EPOCHS} 轮")
print("-" * 40)

train_loss, train_acc = train_model()
test_loss, test_acc = test_model()

print(f"\n本轮总结:")
print(f"训练损失: {train_loss:.4f}, 训练准确率: {train_acc:.2f}%")
print(f"测试损失: {test_loss:.4f}, 测试准确率: {test_acc:.2f}%")

if test_acc > best_test_acc:
best_test_acc = test_acc
torch.save(model.state_dict(), 'mnist_best_model.pth')
print(f"保存最佳模型，测试准确率: {best_test_acc:.2f}%")

print(f"\n训练完成！最佳测试准确率: {best_test_acc:.2f}%")
print("3002") # 新增：训练完成后再次输出3002

def predict_single_image(image, label):
model.eval()
with torch.no_grad():
image = image.unsqueeze(0).to(device)
output = model(image)
_, predicted = torch.max(output, 1)

image = image.squeeze().cpu().numpy()

plt.figure(figsize=(4,4))
plt.imshow(image, cmap='gray')
plt.title(f"真实标签: {label}, 预测标签: {predicted.item()}")
plt.axis('off')
plt.show()

sample_image, sample_label = test_dataset[10]
predict_single_image(sample_image, sample_label)