PyTorch模型性能评估实战：MNIST分类的三大关键问题解析

在利用 PyTorch 构建和训练机器学习模型时，开发者常会遇到一些看似基础但至关重要的问题。本文将以经典的 MNIST 手写数字分类任务为背景，深入探讨其中三个典型问题的解决方法，涵盖代码结构、设备管理与性能评估指标。

核心问题解析与实现

1. 模块导入与脚本执行：if __name__ == '__main__': 的作用

在 Python 脚本中，if __name__ == '__main__': 是一个常见的代码守卫，它用于区分模块是被导入还是被直接运行。这个机制对于编写可复用的代码至关重要，尤其是在进行 Python 项目开发和软件测试时。

• 模块重用：当你编写一个包含函数和类的 .py 文件（模块），并希望在其他脚本中导入它们时，你不希望导入操作触发模块中的训练或测试代码。if __name__ == ‘__main__‘: 内部的代码只在直接运行该文件时执行，确保了模块的纯净性。
• 单元测试：进行自动化测试时，测试框架会导入你的模块。如果没有这个守卫，主程序代码（如训练循环）可能会意外执行，干扰测试流程。
• 避免副作用：某些初始化或配置代码只在作为独立程序运行时才有意义，该守卫可以有效避免在导入时产生不必要的副作用（如下载数据、分配大量资源）。

在我们的 MNIST 训练脚本中，它将所有模型定义、数据加载、训练和测试逻辑包裹起来，确保这部分代码仅在直接运行此脚本时启动。

2. 设备无关的代码：张量如何转移到 CPU？

在 PyTorch 中，无论张量（Tensor）的形状如何（无论是标量、向量还是多维矩阵），将其在 CPU 和 GPU 之间转移的方法是统一的。使用 .to() 方法是实现设备转移的标准且安全的方式。

import torch

# 创建一个形状为 [3, 3] 的张量，可能位于 GPU 上
tensor_gpu = torch.randn(3, 3, device=‘cuda:0‘)

# 使用 .to() 方法将其转移到 CPU，此方法适用于任何形状的张量
tensor_cpu = tensor_gpu.to(‘cpu‘)

print(f‘转移后的设备: {tensor_cpu.device}‘)
# 输出: transfered device: cpu

在模型的测试阶段，我们计算出的准确率 acc 是一个 Python 浮点数（由 .item() 从标量张量转换而来），它本身不涉及设备问题。但如果需要处理非标量的中间张量结果，.to(‘cpu’) 是通用的解决方案。

3. 超越准确率：精确度、召回率与 F1 分数的实现

准确率（Accuracy）是分类任务中最直观的指标，但在类别不平衡的数据集中，它可能具有误导性。为了更全面地评估 人工智能 模型性能，我们通常需要结合精确度（Precision）、召回率（Recall）和 F1 分数（F1-Score）进行综合分析。

• 精确度：在所有被模型预测为正类的样本中，真正为正类的比例。它衡量了模型预测结果的“精确性”。
• 召回率：在所有真实为正类的样本中，被模型正确预测为正类的比例。它衡量了模型对正类的“识别覆盖率”。
• F1 分数：精确度和召回率的调和平均数，旨在找到一个平衡点，在两者之间取得权衡。

我们可以使用 sklearn.metrics 库方便地计算这些指标：

from sklearn.metrics import precision_score, recall_score, f1_score

# 示例：真实标签和模型预测标签
true_labels = [1, 0, 1, 1, 0, 0]
predicted_labels = [1, 1, 1, 0, 0, 1]

# 计算指标， ‘weighted‘ 平均方式考虑了类别不平衡
precision = precision_score(true_labels, predicted_labels, average=‘weighted‘)
recall = recall_score(true_labels, predicted_labels, average=‘weighted‘)
f1 = f1_score(true_labels, predicted_labels, average=‘weighted‘)

print(f‘精确度 (Precision): {precision:.4f}‘)
print(f‘召回率 (Recall): {recall:.4f}‘)
print(f‘F1 分数 (F1-Score): {f1:.4f}‘)

要将这些指标集成到 PyTorch 的测试循环中，你需要在循环中累积所有批次的预测结果和真实标签，然后在循环结束后统一计算。

完整实战示例：MNIST 分类模型

以下代码整合了上述要点，构建了一个完整的训练和测试流程。

import torch
import matplotlib.pyplot as plt
from torch.optim import SGD
from torch import nn
from torchvision import datasets
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader

# 使用 if __name__ == '__main__' 守卫主执行逻辑
if __name__ == '__main__':
    # 设备配置
    device = ‘cuda‘ if torch.cuda.is_available() else ‘cpu‘
    print(f‘使用的设备: {device}‘)

    # 加载 MNIST 数据集
    train_ds = datasets.MNIST(
        root=‘data‘,
        train=True,
        download=True,
        transform=ToTensor(),
    )
    test_ds = datasets.MNIST(
        root=‘data‘,
        train=False,
        download=True,
        transform=ToTensor(),
    )

    batch_size = 256
    train_dataloader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4)
    test_dataloader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=4)

    # 定义一个简单的全连接网络
    class MyNet(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc1 = nn.Linear(28 * 28, 512)
            self.fc2 = nn.Linear(512, 10)

        def forward(self, x):
            x = torch.flatten(x, start_dim=1)  # 将图像展平为向量
            x = self.fc1(x)
            out = self.fc2(x)
            return out

    # 初始化模型、损失函数和优化器
    net = MyNet().to(device)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = SGD(net.parameters(), lr=1e-3)

    losses = []  # 记录损失

    # 训练函数
    def train(dataloader, model, loss_fn, optimizer):
        model.train()
        total_batches = len(dataloader)
        for batch_idx, (X, y) in enumerate(dataloader):
            X, y = X.to(device), y.to(device)
            pred = model(X)
            loss = loss_fn(pred, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            losses.append(loss.item())
            if batch_idx % 100 == 0:
                print(f‘[{batch_idx+1:>4d}/{total_batches}] 损失: {loss.item():.4f}‘)

    # 测试函数（目前仅计算准确率）
    def test(dataloader, model, loss_fn):
        model.eval()
        size = len(dataloader.dataset)
        correct = 0
        with torch.no_grad():
            for X, y in dataloader:
                X, y = X.to(device), y.to(device)
                pred = model(X)
                correct += (pred.argmax(1) == y).type(torch.float).sum().item()
        acc = correct / size
        print(f‘测试准确率: {acc:.4f}\n‘)

    # 执行训练和测试
    epochs = 5
    for epoch in range(epochs):
        print(f‘Epoch {epoch+1}/{epochs}‘)
        train(train_dataloader, net, loss_fn, optimizer)
        test(test_dataloader, net, loss_fn)

    # 可选：绘制训练损失曲线
    plt.plot(losses)
    plt.xlabel(‘迭代次数‘)
    plt.ylabel(‘损失‘)
    plt.title(‘训练损失变化‘)
    plt.show()

运行上述代码，你将在完成训练后看到类似以下的输出，并生成损失曲线图：

使用的设备: cuda
Epoch 1/5
[   1/235] 损失: 2.3025
[ 101/235] 损失: 1.9233
[ 201/235] 损失: 1.1660
测试准确率: 0.8243
...

总结

通过本次实践，我们系统地解决了 PyTorch 模型开发中的三个关键环节：首先，利用 if __name__ == '__main__': 构建了符合工程规范的脚本结构，确保了代码的可复用性与可测试性；其次，掌握了使用 .to() 方法实现设备无关的编码，使模型能灵活运行于不同硬件环境；最后，在准确率之外，探讨了精确度、召回率和 F1 分数等更细致的模型评估指标及其计算方法，为全面优化模型性能打下了坚实基础。将这些知识融会贯通，能显著提升机器学习项目的开发效率和模型质量。