AI-test-mnist-demo-BP-win.zip window下，mnist训练手写体分类的代码样例

import torch import sys import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms import matplotlib.pyplot as plt # 检查设备 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"使用设备: {device}") # 1. 设置数据集的路径 data_dir = "D:\py_work\mnist_ai_bp\AI_test_mnist_demo_BP" # 替换为您的实际路径 # 2. 数据预处理 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) # 使用MNIST的均值和标准差 ]) # 3. 加载数据集 trainset = torchvision.datasets.MNIST(root=data_dir, train=True, download=True, transform=transform) testset = torchvision.datasets.MNIST(root=data_dir, train=False, download=True, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True) testloader = torch.utils.data.DataLoader(testset, batch_size=1000, shuffle=False) # 可视化部分训练数据 examples = enumerate(trainloader) batch_idx, (example_data, example_targets) = next(examples) fig = plt.figure() for i in range(6): plt.subplot(2,3,i+1) plt.tight_layout() plt.imshow(example_data[i][0], cmap='gray', interpolation='none') plt.title(f"标签: {example_targets[i].item()}") plt.xticks([]) plt.yticks([]) plt.show() # 4. 定义神经网络模型 class BPnetwork(nn.Module): def __init__(self): super(BPnetwork, self).__init__() self.linear1 = nn.Linear(28 * 28, 128) self.ReLU1 = nn.ReLU() self.linear2 = nn.Linear(128, 64) self.ReLU2 = nn.ReLU() self.linear3 = nn.Linear(64, 10) self.log_softmax = nn.LogSoftmax(dim=1) def forward(self, x): x = x.view(x.size(0), -1) # 将输入数据展平成一维 x = self.linear1(x) x = self.ReLU1(x) x = self.linear2(x) x = self.ReLU2(x) x = self.linear3(x) x = self.log_softmax(x) return x # 5. 创建模型、损失函数和优化器 model = BPnetwork().to(device) criterion = nn.NLLLoss() # 负对数似然损失函数 optimizer = optim.Adam(model.parameters(), lr=0.001) # 使用Adam优化器 # 6. 模型训练 epochs = 10 for epoch in range(epochs): model.train() # 设置模型为训练模式 running_loss = 0.0 for images, labels in trainloader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() # 清空梯度 output = model(images) # 前向传播 loss = criterion(output, labels) # 计算损失 loss.backward() # 反向传播 optimizer.step() # 更新权重 running_loss += loss.item() average_loss = running_loss / len(trainloader) print(f"Epoch {epoch+1}, Loss: {average_loss:.4f}") # 在每个epoch结束后，评估模型在测试集上的性能 model.eval() # 设置模型为评估模式 correct = 0 total = 0 with torch.no_grad(): for images, labels in testloader: images, labels = images.to(device), labels.to(device) output = model(images) _, predicted = torch.max(output, 1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = 100 * correct / total print(f"测试准确率: {accuracy:.2f}%") # 7. 保存模型 torch.save(model.state_dict(), 'mnist_bpnetwork.pth') print("模型已保存为 mnist_bpnetwork.pth") # 8. 可视化部分测试结果 model.eval() examples = enumerate(testloader) batch_idx, (example_data, example_targets) = next(examples) example_data, example_targets = example_data.to(device), example_targets.to(device) with torch.no_grad(): output = model(example_data) fig = plt.figure() for i in range(6): plt.subplot(2,3,i+1) plt.tight_layout() plt.imshow(example_data[i].cpu()[0], cmap='gray', interpolation='none') pred = output.data.max(1, keepdim=True)[1][i].item() plt.title(f"预测: {pred}") plt.xticks([]) plt.yticks([]) plt.show()