Pytorch：读取手写字体识别数据集，实现AlexNet训练、测试与模型保存代码全解析

```python import torch import torchvision import torchvision.transforms as transforms import torch.nn as nn import torch.optim as optim import torchvision.models as models # 加载数据集 transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2) testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform) testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2) classes = tuple(str(i) for i in range(10)) # 定义网络 net = models.alexnet() # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # 训练网络 for epoch in range(2): # 多批次循环 running_loss = 0.0 for i, data in enumerate(trainloader, 0): # 获取输入 inputs, labels = data # 梯度清零 optimizer.zero_grad() # 正向传播，反向传播，优化 outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() # 打印统计信息 running_loss += loss.item() if i % 2000 == 1999: # 每2000批次打印一次 print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000)) running_loss = 0.0 print('Finished Training') # 测试网络 correct = 0 total = 0 with torch.no_grad(): for data in testloader: images, labels = data outputs = net(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print('Accuracy of the network on the 10000 test images: %d %%' % ( 100 * correct / total)) # 保存模型 PATH = './mnist_alexnet.pth' torch.save(net.state_dict(), PATH) ``` ######[AI写代码神器 | 1466点数解答 | 2024-10-23 23:03:28]