写

使用 MNIST 数据集，训练一个识别手写数字的模型

导入必要的库

import torch
import torch.nn as nn # 用于构建神经网络
import torch.nn.functional as F # 各种函数，比如激活函数
import torch.optim as optim # 优化算法，比如 Adam
from torchvision import datasets, transforms # 视觉向量化
import torch.utils
import torch.utils.data

定义数据类型 & 加载数据集

这段代码定义了一个数据转换流程（transform），目的是将输入的图像数据进行两步处理：

  • 转换为Tensor：把图像数据变成一种适合深度学习模型处理的格式
  • 标准化：调整图像数据的数值范围，让它的分布更适合模型训练

transforms.Compose 定义了一个标准化流程，将多个转换步骤结合在一起

# 转换图片类型为 Tensor，并且标准化，这里的标准化数值时针对 MNIST 数据集的，均值为0.1307，标准差为0.3081
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081))
])

下面加载数据集，由于 MNIST 数据集是在网上的，所以需要下载，这里使用 datasets.MNIST('./MNIST', train=True, download=True, transform=transform), 来指定下载，其中 ./MNIST 是指定了数据集的存储路径，train= 是指定了是作为训练集还是测试集，transform 是选择上面的预处理转换流程，将图片类型转换为 Tensor，并且进行标准化

torch.utils.data.DataLoader(...) 是 PyTorch 给的一个加载数据集的工具，dataset 传入的上面的 MNIST 数据集，batch_size 是每次传入模型的数据，shuffle 是是否打乱数据集，我们这里对训练集进行打乱，对测试集不打乱

# 下载并加载训练集
train_loader = torch.utils.data.DataLoader(
    dataset=datasets.MNIST('./MNIST', train=True, download=True, transform=transform),
    batch_size=64, shuffle=True
)

# 下载并且加载测试集
test_loader = torch.utils.data.DataLoader(
    dataset=datasets.MNIST('./MNIST', train=False, download=True, transform=transform),
    batch_size=1000, shuffle=False
)

定义神经网络结构

网络结构

这个神经网络有两层全连接层（nn.Linear），也就是两个线性变换层：

  • self.fc1 = nn.Linear(28*28, 128)：第一层（输入层到隐藏层）。输入是 28*28=784（因为 MNIST 图片是 28 x 28 像素，展平后是一个长度为 784 的向量），输出是 128 个神经元。这就像从 784 个输入点连接到 128 个中间点
  • self.fc2 = nn.Linear(128, 10)：第二层（隐藏层到输出层）。输入是 128 个神经元，输出是 10 个神经元，分别对应 0 到 9 这 10 个数字的分类分数
输入 (784维) --> 全连接层1 (784到128) --> ReLU激活 --> 全连接层2 (128到10) --> 输出 (10维得分)

前向传播（定义数据从输入到输出）

  • x = x.view(-1, 28*28)：将数据进行降维，x 是从上面的 dataLoader 拿到的数据，这时候的数据为 tensor 类型，根据 MNIST 数据集的定义，其中的 tensor 数据形状应该为 [batch_size, 1, 28, 28]，这时候使用 .view 方法对其进行降维，将其变成一维数据，即 28*28，其中 -1 表示自动计算维度大小
  • x = F.relu(self.fc1(x))，这里对第一层神经网络的线性输出使用了 ReLU 激活函数
  • x = self.fc2(x) ，最后进行输出
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 128) # 输入层到隐藏层
        self.fc2 = nn.Linear(128, 10) # 隐藏层到输出层，10个数字，故10个神经元

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = F.relu(self.fc1(x))
        x= self.fc2(x)
        return x

设置模型和优化器

对于使用 Apple silicon 或者 AMD GPUs 的 Mac 设备来说，可以使用 Metal Performance Shaders (MPS) 来代替 CUDA 进行 GPU 高性能训练

if torch.backends.mps.is_available():
    device = torch.device("mps")
    print("使用 MPS 后端")
else:
    device = torch.device("cpu")
    print("使用 CPU")
print(f"当前设备：{device}")

model = Net().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)

定义训练函数

这里的 train 函数需要传入 5 个形参，之后设置模式为训练模式 model.train()

for batch_idx, (data, target) in enumerate(train_loader): 这个循环会分别从 train_loader 中拿出 batch 的数据以及编号

data, target = data.to(device), target.to(device) 会将数据移动到对应的设备中

optimizer.zero_grad() 对前面的梯度进行清空，防止影响这一批的训练

loss = F.cross_entropy(output=output, target=target) 使用交叉熵损失函数，进行计算，并且使用 loss.backward() 对损失函数进行反向传播

optimizer.step() 根据反向传播计算出来的梯度对其进行参数更新

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()

        if batch_idx % 100 == 0:
            print(f'Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] Loss: {loss.item():.6f}')

测试函数

这里的 test 函数需要传入三个形参，model 是要评估的模型，device 是使用计算的设备，test_loader 是传入的测试集

torch.no_grad() 是设置不需要使用梯度

这里是将数据从测试集中拿出来，并且将其放入设备中，准备计算

for data, target in test_loader:
            data, target = data.to(device), target.to(device)

output = model(data) 是开始使用模型进行预测

test_loss += F.cross_entropy(output, target, reduction='sum').item() 使用交叉熵损失函数来计算预测和实际之间的差值，.item() 将 PyTorch 中的张量转化为普通的 Python 变量，以便于累加

pred = output.argmax(dim=1) 选择其中概率最高的作为预测结果

correct += pred.eq(target).sum().item() 将预测的和目标结果进行比较，返回一个布尔张量，并且将其累加，最终转化成 Python 数值，进行统计预测成功的

最后进行统计输出

test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)

    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')
def test(model, device, test_loader):
    model.eval() # 模型设置为评估模型
    test_loss = 0 # 设置错误累计
    corrent = 0  # 设置成功

    # 在测试模式下不需要梯度
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)

            output = model(data)
            test_loss += F.cross_entropy(output, target, reduction='sum').item()
            pred = output.argmax(dim=1)
            corrent += pred.eq(target).sum().item()

    test_loss /= len(test_loader.dataset)
    accuracy = 100. * corrent / len(test_loader.dataset)

    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {corrent}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')

开始训练

这里选择迭代 6 次，在训练集训练结束后直接使用测试集进行检验

import time
all_time = time.time()

for epoch in range(1, 6):
    start_time = time.time()  # 记录开始时间
    train(model, device, train_loader, optimizer, epoch)
    test(model, device, test_loader)
    elapsed = time.time() - start_time  # 计算当前 epoch 用时
    print(f"Epoch {epoch} 用时: {elapsed:.2f} 秒")

spend_time = time.time() - all_time
print(f'\n总耗时：{spend_time}')

训练结束后，可以选择将训练好的模型（参数）保存在文件中：

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

随机打开一张图检验

import matplotlib.pyplot as plt
import numpy as np

# 获取一张图片
examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)

# 取第0张
example = example_data[0].unsqueeze(0).to(device)
output = model(example)
pred = output.argmax(dim=1)

# 显示图像
plt.imshow(example.cpu().squeeze(), cmap='gray')
plt.title(f'pred result: {pred.item()}')
plt.axis('off')
plt.show()

使用 CPU 或者 GPU 训练的时长区别

笔者所使用的是 Macbook Air M2 16G 版本，在这里我们分别测试一下使用 Metal 和 CPU 训练的时间对比

训练的代码是这个，在这里我已经先定义好了

import time
all_time = time.time()

for epoch in range(1, 6):
    start_time = time.time()  # 记录开始时间
    train(model, device, train_loader, optimizer, epoch)
    test(model, device, test_loader)
    elapsed = time.time() - start_time  # 计算当前 epoch 用时
    print(f"Epoch {epoch} 用时: {elapsed:.2f} 秒")

spend_time = time.time() - all_time
print(f'\n总耗时：{spend_time}')

使用 CPU

Epoch: 1 [0/60000] Loss: 2.277056
Epoch: 1 [6400/60000] Loss: 0.644652
Epoch: 1 [12800/60000] Loss: 0.222628
Epoch: 1 [19200/60000] Loss: 0.176244
Epoch: 1 [25600/60000] Loss: 0.378818
Epoch: 1 [32000/60000] Loss: 0.290092
Epoch: 1 [38400/60000] Loss: 0.162229
Epoch: 1 [44800/60000] Loss: 0.139718
Epoch: 1 [51200/60000] Loss: 0.346534
Epoch: 1 [57600/60000] Loss: 0.289447

Test set: Average loss: 0.1444, Accuracy: 9545/10000 (95.45%)

Epoch 1 用时: 3.59 秒
Epoch: 2 [0/60000] Loss: 0.144912
Epoch: 2 [6400/60000] Loss: 0.072045
Epoch: 2 [12800/60000] Loss: 0.135150
Epoch: 2 [19200/60000] Loss: 0.104258
Epoch: 2 [25600/60000] Loss: 0.049891
Epoch: 2 [32000/60000] Loss: 0.301743
Epoch: 2 [38400/60000] Loss: 0.177394
Epoch: 2 [44800/60000] Loss: 0.052172
Epoch: 2 [51200/60000] Loss: 0.280824
Epoch: 2 [57600/60000] Loss: 0.246844
...

Epoch 5 用时: 3.49 秒

总耗时：17.6266610622406

使用 GPU

Epoch: 1 [0/60000] Loss: 2.377831
Epoch: 1 [6400/60000] Loss: 0.408726
Epoch: 1 [12800/60000] Loss: 0.485179
Epoch: 1 [19200/60000] Loss: 0.278891
Epoch: 1 [25600/60000] Loss: 0.307072
Epoch: 1 [32000/60000] Loss: 0.113186
Epoch: 1 [38400/60000] Loss: 0.199770
Epoch: 1 [44800/60000] Loss: 0.225306
Epoch: 1 [51200/60000] Loss: 0.089396
Epoch: 1 [57600/60000] Loss: 0.151413

Test set: Average loss: 0.1318, Accuracy: 9604/10000 (96.04%)

Epoch 1 用时: 4.77 秒
Epoch: 2 [0/60000] Loss: 0.113682
Epoch: 2 [6400/60000] Loss: 0.044853
Epoch: 2 [12800/60000] Loss: 0.142270
Epoch: 2 [19200/60000] Loss: 0.122028
Epoch: 2 [25600/60000] Loss: 0.026501
Epoch: 2 [32000/60000] Loss: 0.095724
Epoch: 2 [38400/60000] Loss: 0.085834
Epoch: 2 [44800/60000] Loss: 0.090036
Epoch: 2 [51200/60000] Loss: 0.242773
Epoch: 2 [57600/60000] Loss: 0.036999

Test set: Average loss: 0.0993, Accuracy: 9718/10000 (97.18%)

Epoch 2 用时: 4.82 秒
Epoch: 3 [0/60000] Loss: 0.029983
Epoch: 3 [6400/60000] Loss: 0.091785
Epoch: 3 [12800/60000] Loss: 0.020735
Epoch: 3 [19200/60000] Loss: 0.058402
Epoch: 3 [25600/60000] Loss: 0.018810
Epoch: 3 [32000/60000] Loss: 0.218197
Epoch: 3 [38400/60000] Loss: 0.098756
Epoch: 3 [44800/60000] Loss: 0.098376
Epoch: 3 [51200/60000] Loss: 0.157909
Epoch: 3 [57600/60000] Loss: 0.023484

Test set: Average loss: 0.0871, Accuracy: 9737/10000 (97.37%)

Epoch 3 用时: 4.76 秒
Epoch: 4 [0/60000] Loss: 0.079279
Epoch: 4 [6400/60000] Loss: 0.013454
Epoch: 4 [12800/60000] Loss: 0.048985
Epoch: 4 [19200/60000] Loss: 0.022868
Epoch: 4 [25600/60000] Loss: 0.061240
Epoch: 4 [32000/60000] Loss: 0.007518
Epoch: 4 [38400/60000] Loss: 0.102460
Epoch: 4 [44800/60000] Loss: 0.040929
Epoch: 4 [51200/60000] Loss: 0.028633
Epoch: 4 [57600/60000] Loss: 0.163401

Test set: Average loss: 0.0762, Accuracy: 9758/10000 (97.58%)

Epoch 4 用时: 4.80 秒
Epoch: 5 [0/60000] Loss: 0.067033
Epoch: 5 [6400/60000] Loss: 0.032746
Epoch: 5 [12800/60000] Loss: 0.014985
Epoch: 5 [19200/60000] Loss: 0.072542
Epoch: 5 [25600/60000] Loss: 0.061535
Epoch: 5 [32000/60000] Loss: 0.034278
Epoch: 5 [38400/60000] Loss: 0.054543
Epoch: 5 [44800/60000] Loss: 0.005301
Epoch: 5 [51200/60000] Loss: 0.049687
Epoch: 5 [57600/60000] Loss: 0.048477

Test set: Average loss: 0.0788, Accuracy: 9755/10000 (97.55%)

Epoch 5 用时: 4.85 秒

总耗时：23.998369932174683

这里算下来，使用 CPU 推理的效果还不如 GPU，后来发现问题出在该数据集的 beach_size 上，beach_size 太小了，没办法发挥出 GPU 并行计算的优势

当笔者将 beach_size 提升到 128 后，效率有明显提升：

Epoch: 1 [0/60000] Loss: 0.002459
Epoch: 1 [12800/60000] Loss: 0.000959
Epoch: 1 [25600/60000] Loss: 0.001181
Epoch: 1 [38400/60000] Loss: 0.001985
Epoch: 1 [51200/60000] Loss: 0.002780

Test set: Average loss: 0.0803, Accuracy: 9799/10000 (97.99%)

Epoch 1 用时: 3.87 秒
Epoch: 2 [0/60000] Loss: 0.006955
Epoch: 2 [12800/60000] Loss: 0.002206
Epoch: 2 [25600/60000] Loss: 0.001883
Epoch: 2 [38400/60000] Loss: 0.003889
Epoch: 2 [51200/60000] Loss: 0.003672

Test set: Average loss: 0.1011, Accuracy: 9759/10000 (97.59%)

Epoch 2 用时: 3.64 秒
Epoch: 3 [0/60000] Loss: 0.005289
Epoch: 3 [12800/60000] Loss: 0.012927
Epoch: 3 [25600/60000] Loss: 0.028744
Epoch: 3 [38400/60000] Loss: 0.016574
Epoch: 3 [51200/60000] Loss: 0.019915

Test set: Average loss: 0.0929, Accuracy: 9788/10000 (97.88%)

Epoch 3 用时: 3.69 秒
Epoch: 4 [0/60000] Loss: 0.001103
Epoch: 4 [12800/60000] Loss: 0.002012
Epoch: 4 [25600/60000] Loss: 0.000979
Epoch: 4 [38400/60000] Loss: 0.009304
Epoch: 4 [51200/60000] Loss: 0.000931

Test set: Average loss: 0.0955, Accuracy: 9782/10000 (97.82%)

Epoch 4 用时: 3.64 秒
Epoch: 5 [0/60000] Loss: 0.000375
Epoch: 5 [12800/60000] Loss: 0.016383
Epoch: 5 [25600/60000] Loss: 0.003329
Epoch: 5 [38400/60000] Loss: 0.004040
Epoch: 5 [51200/60000] Loss: 0.012434

Test set: Average loss: 0.0964, Accuracy: 9784/10000 (97.84%)

Epoch 5 用时: 3.75 秒

总耗时：18.60126280784607
最后修改：2025 年 04 月 26 日
© 允许规范转载
如果觉得我的文章对你有用，请随意赞赏