[Pytorch案例实践001]手写minist数字识别

一、流程介绍        

该案例实现了使用PyTorch框架训练一个卷积神经网络（CNN）来识别MNIST手写数字数据集中的数字。下面是对整个流程的总结：

1. 环境设置：
   导入必要的库，如 `torch`, `torch.nn`, `torch.optim`, `torchvision.datasets`, `torchvision.transforms`, `torch.utils.data`, `matplotlib.pyplot` 和 `tqdm`。
   设定运行设备为 GPU（如果有可用的 GPU）或 CPU。

2. 数据预处理：
   定义数据转换步骤，包括将图像转为张量和标准化。
   加载MNIST数据集的训练集和测试集，并应用定义好的数据转换。
   使用 `DataLoader` 对数据集进行批处理和打乱顺序。

3. 模型定义：
   构建一个简单的CNN模型，包含两层卷积层，每层后跟着ReLU激活函数和最大池化层，以及几个全连接层。
   将模型放到选定的设备上。

4. 损失函数与优化器：
   选择交叉熵损失作为分类任务的标准损失函数。
   使用 Adam 优化器更新模型参数，设定学习率为 0.001。

5. 训练和测试函数：
   定义训练函数，其中包括了模型训练的循环逻辑，如前向传播、反向传播和参数更新。
   定义测试函数，用于评估模型在测试集上的表现。

6. 主训练循环：
   进行多个周期（epochs）的训练和测试。
   每个周期结束后记录训练损失、训练准确率、测试损失和测试准确率。
   如果当前周期的测试准确率高于之前的所有周期，则保存模型的状态字典。

7. 结果可视化：
   使用 Matplotlib 绘制训练和测试损失随周期变化的曲线，以及训练和测试准确率随周期变化的曲线。保存并显示这些曲线图。

二、完整带注释代码

代码附录如下：

import torch
import torch.nn as nn  # 导入神经网络模块
import torch.optim as optim  # 导入优化器模块
from torchvision import datasets, transforms  # 导入数据集和预处理模块
from torch.utils.data import DataLoader  # 导入数据加载模块
import matplotlib.pyplot as plt  # 导入绘图模块
from tqdm import tqdm  # 导入进度条模块

# 定义设备，GPU如果可用则使用，否则使用CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据预处理，将图像转换为张量并进行标准化
# 用于将MNIST数据集标准化，其中 0.1307 是训练集的均值，0.3081 是训练集的标准差。
# 标准化操作是对每个像素值进行变换，使数据具有零均值和单位方差。这对加速训练收敛和提高模型性能有帮助
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 下载并加载MNIST训练和测试数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=False, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=False, transform=transform)

# 使用DataLoader加载数据集
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=1000, shuffle=False)

# 使用nn.Sequential构建模型，包括卷积层、激活层、池化层和全连接层
model = nn.Sequential(
    nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),  # 第一个卷积层，输入通道1，输出通道32，卷积核大小3x3
    nn.ReLU(),  # ReLU激活函数
    nn.MaxPool2d(kernel_size=2, stride=2),  # 最大池化层，池化核大小2x2
    nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),  # 第二个卷积层，输入通道32，输出通道64，卷积核大小3x3
    nn.ReLU(),  # ReLU激活函数
    nn.MaxPool2d(kernel_size=2, stride=2),  # 最大池化层，池化核大小2x2
    nn.Flatten(),  # 展平操作，将多维张量展平成一维
    nn.Linear(64 * 7 * 7, 128),  # 全连接层，输入尺寸64*7*7，输出尺寸128
    nn.ReLU(),  # ReLU激活函数
    nn.Linear(128, 64),  # 全连接层，输入尺寸128，输出尺寸64
    nn.ReLU(),  # ReLU激活函数
    nn.Linear(64, 10)  # 全连接层，输入尺寸64，输出尺寸10（对应10个类别）
).to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 使用交叉熵损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)  # 使用Adam优化器，学习率为0.001


# 训练函数
def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()  # 设置模型为训练模式
    train_loss = 0
    correct = 0
    progress_bar = tqdm(train_loader)  # 创建进度条
    for data, target in progress_bar:
        data, target = data.to(device), target.to(device)  # 将数据和目标移动到设备上
        optimizer.zero_grad()  # 清空梯度
        output = model(data)  # 前向传播
        loss = criterion(output, target)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数
        train_loss += loss.item()  # 累积训练损失

        # 对于每个样本，从 output 中找出预测概率最高的类别的索引，并将其保存在 pred 中。
        # 通过使用 keepdim=True，保持输出的形状为 (batch_size, 1)，使其更易于与目标标签进行比较。
        # output: 这是神经网络的前向传播结果。对于分类任务来说，output 的形状通常是 (batch_size, num_classes)，
        # 其中 batch_size 是当前批次的样本数量，num_classes 是类别数量（在MNIST数据集中，num_classes 是10，因为有10个数字类别）。
        # argmax(dim=1): 这个函数用于返回指定维度上最大值的索引。在这里，dim=1 表示我们在类别维度上进行操作。也就是说，对于每个样本，
        # 我们找出在10个类别中预测概率最高的那个类别的索引。例如，如果某个样本的输出是 [0.1, 0.3, 0.1, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]，
        # argmax(dim=1) 将返回 3，因为 0.5 是最大值，索引是 3。
        # keepdim=True: 这个参数表示保持输出张量的维度。具体来说，如果输入张量的形状是 (batch_size, num_classes)，
        # 使用 argmax(dim=1, keepdim=True) 后输出张量的形状将是 (batch_size, 1)。这样可以方便后续的操作，
        # 例如计算预测的正确性。如果 keepdim=False（默认值），输出张量的形状将是 (batch_size,)。
        pred = output.argmax(dim=1, keepdim=True)  # 获取预测结果

        # pred: 这是模型预测的结果，形状为 (batch_size, 1)。每个元素是模型对相应样本的预测类别索引。
        # target: 这是实际的目标标签，形状为 (batch_size,)。每个元素是样本的实际类别索引。
        # target.view_as(pred): 这部分代码将 target 视图转换为与 pred 相同的形状 (batch_size, 1)。
        # view_as 是 PyTorch 中的一种方法，它可以将一个张量的形状调整为另一个张量的形状。
        # pred.eq(target.view_as(pred)): 这部分代码比较 pred 和 target 是否相等，返回一个布尔张量。
        # 对于每个元素，如果 pred 和 target 相等，则返回 True，否则返回 False。形状仍然是 (batch_size, 1)。
        # .sum(): 这部分代码对布尔张量进行求和。由于 True 在数值上等于 1，False 等于 0，因此 sum() 返回正确预测的数量。
        # .item(): 这部分代码将单个标量张量转换为 Python 数值。 sum() 的结果是一个标量张量，而 .item() 则提取其中的数值。
        # correct += ...: 将当前批次中正确预测的数量累加到 correct 变量中。
        correct += pred.eq(target.view_as(pred)).sum().item()  # 计算正确预测的数量
        progress_bar.set_description(f'Epoch {epoch}, Loss: {loss.item():.4f}')  # 更新进度条描述
    train_loss /= len(train_loader)  # 计算平均训练损失
    train_accuracy = correct / len(train_loader.dataset)  # 计算训练集准确率
    return train_loss, train_accuracy


# 测试函数
def test(model, device, test_loader, criterion):
    model.eval()  # 设置模型为评估模式
    test_loss = 0
    correct = 0
    with torch.no_grad():  # 不计算梯度
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)  # 将数据和目标移动到设备上
            output = model(data)  # 前向传播
            test_loss += criterion(output, target).item()  # 累积测试损失
            pred = output.argmax(dim=1, keepdim=True)  # 获取预测结果
            correct += pred.eq(target.view_as(pred)).sum().item()  # 计算正确预测的数量
    test_loss /= len(test_loader)  # 计算平均测试损失
    test_accuracy = correct / len(test_loader.dataset)  # 计算测试集准确率
    return test_loss, test_accuracy


# 主训练循环
num_epochs = 10  # 训练周期数
train_losses = []  # 存储训练损失
train_accuracies = []  # 存储训练准确率
test_losses = []  # 存储测试损失
test_accuracies = []  # 存储测试准确率

best_accuracy = 0.0  # 初始化最佳准确率

for epoch in range(1, num_epochs + 1):
    # 训练模型
    train_loss, train_accuracy = train(model, device, train_loader, optimizer, criterion, epoch)
    # 测试模型
    test_loss, test_accuracy = test(model, device, test_loader, criterion)

    # 记录损失和准确率
    train_losses.append(train_loss)
    train_accuracies.append(train_accuracy)
    test_losses.append(test_loss)
    test_accuracies.append(test_accuracy)

    # 打印每个周期的损失和准确率
    print(f'Epoch {epoch}: Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}, '
          f'Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.4f}')

    # 保存性能最好的模型
    if test_accuracy > best_accuracy:
        best_accuracy = test_accuracy
        torch.save(model.state_dict(), 'best_model.pth')

# 绘制损失函数和准确率曲线
plt.figure(figsize=(12, 5))

# 绘制训练和测试损失曲线
plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Train Loss')
plt.plot(test_losses, label='Test Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss Curve')
plt.legend()

# 绘制训练和测试准确率曲线
plt.subplot(1, 2, 2)
plt.plot(train_accuracies, label='Train Accuracy')
plt.plot(test_accuracies, label='Test Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Accuracy Curve')
plt.legend()

plt.tight_layout()
plt.savefig('loss_accuracy_curves.png')  # 保存曲线图
plt.show()  # 显示曲线图

三、总结

        该案例完成了一个完整的机器学习项目流程，包括数据加载、模型构建、训练、验证以及结果可视化。演示了如何使用PyTorch进行端到端的深度学习实验，并且能够针对MNIST数据集训练出一个表现良好的卷积神经网络模型。