写一个基于pytorch的手写数字识别程序

创建一个基于PyTorch的手写数字识别程序通常包括以下几个步骤：

1. **导入库**：首先需要导入必要的库，如`torch`, `torchvision`（用于数据处理），以及一些可视化工具。

import torch
import torchvision
from torchvision import datasets, transforms
import torch.nn as nn
import torch.optim as optim

2. **数据预处理**：加载MNIST数据集并对其进行预处理，这包括归一化像素值。

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

3. **数据加载器**：创建数据迭代器以便于模型训练和验证。

batch_size = 64
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

4. **构建模型**：设计卷积神经网络（CNN）结构，例如LeNet、AlexNet等，作为手写数字识别模型。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 网络层定义...
        
    def forward(self, x):
        # 前向传播过程...

net = Net()

5. **编译模型**：设定损失函数（交叉熵）、优化器（SGD或其他）。

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

6. **训练**：遍历训练集，前向传播、计算损失、反向传播并更新权重。

num_epochs = 10
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 执行前向传播、计算损失和梯度更新
        ...
        print(f"Epoch {epoch+1}, Step {i+1}: Loss {loss.item():.4f}")

7. **评估**：在测试集上评估模型性能。

correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy}%")

8. **保存模型**：如果想要在未来复用模型，可以保存训练好的权重。

torch.save(net.state_dict(), 'mnist_model.pth')