### PyTorch 实现 MNIST 手写数字分类 #### 导入必要的库 为了构建一个简单的神经网络来对手写数字进行分类，首先需要导入一些基础的 Python 库以及 PyTorch 的相关模块。 ```python import torch from torch import nn, optim import torch.nn.functional as F from torchvision import datasets, transforms from torch.utils.data import DataLoader ``` #### 数据预处理与加载 定义数据转换操作并创建 `DataLoader` 来批量读取训练和测试数据集中的图片。这里设置了批大小(batch size)，以便每次迭代可以获取固定数量的数据样本用于训练或评估模型性能[^1]。 ```python transform = transforms.Compose([ transforms.ToTensor(), ]) train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform) batch_size = 64 train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) ``` #### 构建卷积神经网络(CNN)模型 设计了一个简单的 CNN 结构来进行图像特征提取工作。此部分包含了两个卷积层(Conv2d), ReLU 激活函数被用来增加非线性特性[^2]。 ```python class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5) self.conv2 = nn.Conv2d(in_channels=10, out_channels=20, kernel_size=5) self.fc1 = nn.Linear(in_features=320, out_features=50) self.fc2 = nn.Linear(in_features=50, out_features=10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2(x), 2)) x = x.view(-1, 320) x = F.relu(self.fc1(x)) x = self.fc2(x) return F.log_softmax(x, dim=1) ``` #### 定义损失函数与优化器 选择了交叉熵作为损失计算方式，并采用随机梯度下降(SGD)算法更新权重参数以最小化误差[^4]。 ```python network = Net() optimizer = optim.SGD(network.parameters(), lr=0.01, momentum=0.9) criterion = nn.CrossEntropyLoss() ``` #### 训练过程 通过循环遍历整个训练集完成多轮次的学习任务，在每个 epoch 中调整模型参数使得预测更加接近真实标签值。 ```python epochs = 5 for epoch in range(epochs): running_loss = 0.0 for data, target in train_loader: optimizer.zero_grad() output = network(data) loss = criterion(output, target) loss.backward() optimizer.step() running_loss += loss.item() print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}') ``` #### 测试模型准确性 最后一步是在独立验证集中检验所建立模型的效果如何，输出前六个样例及其对应的预测结果可视化展示出来[^5]。 ```python examples = enumerate(test_loader) batch_idx, (example_data, example_targets) = next(examples) with torch.no_grad(): output = network(example_data) _, predicted = torch.max(output, 1) fig = plt.figure(figsize=(8, 3)) for i in range(6): ax = fig.add_subplot(2, 3, i + 1) ax.imshow(example_data[i][0], cmap='gray', interpolation='none') ax.set_title(f'Pred: {predicted[i].item()}') plt.tight_layout() plt.show() ```