今天开始使用 PyTorch 完成一个经典的深度学习项目——MNIST 手写数字识别。MNIST 数据集包含 60000 张训练图片和 10000 张测试图片每张图片都是 28×28 像素的灰度图标签为 0 到 9 的数字。这个项目的主要流程是加载数据转换为 Tensor 使用 DataLoader 分批读取构建神经网络前向传播计算损失反向传播更新参数测试模型准确率一、导入环境并查看版本import torch import torchvision import torchaudio print(torch.__version__) print(torchvision.__version__) print(torchaudio.__version__)如果需要使用代理下载数据集可以设置环境变量import os os.environ[HTTP_PROXY] http://127.0.0.1:7877 os.environ[HTTPS_PROXY] http://127.0.0.1:7877这里需要注意版本属性应该写成__version__前后各有两个下划线。注如果是在Linux环境下或者服务器部署我们也可以从https://download.pytorch.org/whl/这个pytorch官方网站上获取下载名称中包含Linux和whl的文件然后转移到Linux系统里。二、加载 MNIST 数据集PyTorch 提供了封装好的 MNIST 数据集可以直接下载和读取from torchvision import datasets from torchvision.transforms import ToTensor training_data datasets.MNIST( rootdata, trainTrue, downloadTrue, transformToTensor() ) test_data datasets.MNIST( rootdata, trainFalse, downloadTrue, transformToTensor() )参数含义如下1.rootdata数据集保存位置2.trainTrue加载训练集3.trainFalse加载测试集4.downloadTrue本地没有数据时自动下载5.transformToTensor()将图片转换为 Tensor。原始图片不能直接输入神经网络需要先转换为 Tensor。转换之后图像的形状通常为[通道数, 高度, 宽度]MNIST 是灰度图因此通道数为 1单张图片的形状为[1, 28, 28]Tensor 与 NumPy 数组类似都可以保存多维数据但 Tensor 可以直接参与 PyTorch 的自动求导也可以放到 GPU 上运行。三、查看数据集中的图片为了确认数据是否加载正确可以使用 Matplotlib 显示部分训练图片from matplotlib import pyplot as plt figure plt.figure() for i in range(9): #取九个数字 img, label training_data[i] figure.add_subplot(3, 3, i 1) #3*3的画板 plt.title(label) plt.axis(off) plt.imshow(img.squeeze(), cmapgray) #.squeeze去除通道维 plt.show()img的形状是[1, 28, 28]而 Matplotlib 显示灰度图片时通常只需要[28, 28]因此使用squeeze()去掉通道维度。最后显示出以上图像。四、使用 DataLoader 分批读取训练神经网络时通常不会一次性把全部数据输入模型而是将数据划分成多个批次。这样做主要有两个好处一是可以降低单次计算的内存占用避免显存或内存溢出二是每个批次相当于一次“小规模抽样”能让梯度更新更平稳训练过程也更稳定。from torch.utils.data import DataLoader train_dataloader DataLoader( training_data, batch_size64 ) test_dataloader DataLoader( test_data, batch_size64 )这里设置batch_size64表示每次读取 64 张图片。可以查看一个批次的数据形状for x, y in test_dataloader: print(fshape of x: {x.shape}) print(fshape of y: {y.shape}) print(ftype of y: {y.dtype}) break输出结果大致为shape of x: torch.Size([64, 1, 28, 28]) shape of y: torch.Size([64]) type of y: torch.int64其中x是图片数据y是对应的数字标签。五、选择运行设备为了让模型能够在 GPU 上运行可以根据当前环境自动选择设备device torch.device( cuda:0 if torch.cuda.is_available() else cpu) print(device)如果电脑安装了 CUDA 并且 PyTorch 能够识别显卡就会使用 CUDA否则使用 CPU。模型和数据必须放在同一个设备上否则会出现设备不一致的错误。六、构建神经网络本次使用一个简单的全连接神经网络from torch import nn class NeuralNetwork(nn.Module): def __init__(self): super().__init__() self.flatten nn.Flatten() self.hidden1 nn.Linear(28 * 28, 128) self.hidden2 nn.Linear(128, 256) self.output nn.Linear(256, 10) def forward(self, x): x self.flatten(x) x self.hidden1(x) x torch.sigmoid(x) x self.hidden2(x) x torch.sigmoid(x) x self.output(x) return x网络结构可以表示为28×28 图片 → Flatten 展平为 784 个特征 → Linear(784, 128) → Sigmoid → Linear(128, 256)→ Sigmoid → Linear(256, 10) → 输出十个类别分数Flatten()会将[1, 28, 28]展平为长度为 784 的向量。因为全连接层要求输入是一维特征所以需要先进行展平。nn.Linear()表示全连接层。第一层将 784 个像素特征映射为 128 个隐藏特征第二层将 128 个特征映射为 256 个特征最后输出 10 个数字类别对应的分数。输出层有 10 个节点分别对应数字 0 到 9。输出值还不是概率而是每个类别的预测分数通常称为 logits。创建模型并放到指定设备GPUmodel NeuralNetwork().to(device) print(model)七、训练模型训练函数的代码如下def train(dataloader, model, loss_fn, optimizer): model.train() batch_size_num 1 for x, y in dataloader: x x.to(device) y y.to(device) pred model(x) loss loss_fn(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() loss_value loss.item() if batch_size_num % 100 0: print(floss: {loss_value:.4f}) batch_size_num 1训练过程主要包含四个步骤。第一步是前向传播pred model(x)模型接收输入图片并计算得到预测结果。第二步是计算损失loss loss_fn(pred, y)损失函数用于衡量预测结果与真实标签之间的差距。损失越小说明模型预测得越准确。第三步是清空梯度optimizer.zero_grad()PyTorch 默认会累积梯度因此每次更新参数前都要先清空上一轮的梯度。第四步是反向传播并更新参数loss.backward() optimizer.step()loss.backward()根据损失函数自动计算每个参数的梯度optimizer.step()根据梯度更新网络参数。八、测试模型测试阶段不需要更新参数只需要计算模型在测试集上的表现def test(dataloader, model, loss_fn): size len(dataloader.dataset) num_batches len(dataloader) model.eval() test_loss 0 correct 0 with torch.no_grad(): for x, y in dataloader: x x.to(device) y y.to(device) pred model(x) test_loss loss_fn( pred, y ).item() correct ( pred.argmax(1) y ).type(torch.float).sum().item() test_loss / num_batches correct / size print( fTest result:\n fAccuracy: {correct * 100:.0f}%\n fAvg loss: {test_loss:.8f}\n )model.eval()会将模型切换到测试模式。torch.no_grad()表示测试阶段不计算梯度可以减少内存占用并提高运行速度。模型输出的形状为[批次大小, 10]每一行对应一张图片的十个类别分数。使用下面的代码可以得到预测类别pred.argmax(1)这里将预测出来的类别与真实标签做对比预测对为True预测错为False.type(torch.float)这里是将布尔类型的张量转换为浮点类型的argmax(1)表示在类别维度上寻找最大值的位置。例如[1.2, 0.3, 5.8, 0.6, ...]最大值出现在下标 2因此模型预测该图片是数字 2。九、损失函数与优化器本次使用交叉熵损失loss_fn nn.CrossEntropyLoss()交叉熵损失常用于多分类任务。MNIST 有 10 个类别因此非常适合使用这个损失函数。需要注意使用CrossEntropyLoss()时模型最后一层不需要手动添加 Softmax。因为交叉熵损失内部已经包含了相关计算直接输入 logits 即可。优化器使用 Adamoptimizer torch.optim.Adam( model.parameters(), lr0.01 )model.parameters()表示将模型中的所有可训练参数交给优化器管理lr0.01是学习率用于控制每次参数更新的步长。十、完整训练流程设置训练轮数并开始训练epochs 10 for i in range(epochs): print(fepoch [{i 1}/{epochs}]) train( train_dataloader, model, loss_fn, optimizer ) print(Done) test( test_dataloader, model, loss_fn )一个 epoch 表示模型完整遍历一次训练集。训练 10 轮后模型会反复学习训练图片中的像素特征和数字结构。随着训练进行损失一般会逐渐下降测试准确率逐步提高。不过学习率、网络结构、激活函数和训练轮数都会影响最终效果。总结今天通过 MNIST 手写数字识别完整体验了 PyTorch 的基本训练流程。首先使用datasets.MNIST加载数据再通过ToTensor()将图片转换为 Tensor使用DataLoader按批次读取数据。模型部分使用了三个全连接层中间加入 Sigmoid 激活函数将 28×28 的图片转换为 10 个数字类别的预测结果。训练时通过前向传播计算预测值再使用交叉熵损失衡量误差最后通过反向传播和 Adam 优化器更新模型参数。这次代码虽然结构比较简单但已经包含了深度学习项目中最核心的几个环节数据集模型损失函数优化器前向传播反向传播模型评估后续还可以继续尝试修改网络层数、隐藏层节点数、激活函数和学习率观察不同参数对准确率和训练速度的影响。也可以将全连接网络改成卷积神经网络进一步提升 MNIST 图片识别效果。