最近有朋友问我PyTorch怎么入门我每次都推荐先做一个数字识别项目。原因很简单数据集好找任务直观网络代码量不大训练十分钟就能看到效果而且从数据加载到模型评估几乎覆盖了深度学习最常见的全流程。我当时用PyTorch实现手写数字识别时在代码里写满了注释一步一个坑踩过来今天把这些内容整理出来希望能给你一条顺畅的上手路径。这篇分享围绕PyTorch图像识别中的数字识别任务展开从环境准备、数据下载、模型设计、训练循环到测试评估每一段代码都有详细注释和设计理由。不管你是刚接触深度学习的萌新还是想快速巩固PyTorch基础的老手都可以直接照着做然后在这个基础上玩出更多花样。1. 为什么拿数字识别当PyTorch入门项目1.1 数字识别任务的前世今生手写数字识别是计算机视觉里最经典的入门任务之一。一张28×28的灰度图里面是一个0到9的手写数字模型需要判断这张图属于哪一类。任务本身不难理解但它把图像识别的核心流程全串起来了图像怎么读入、怎么转成张量、怎么归一化、怎么喂给网络、怎么得到分类结果。这个任务背后是几十年的算法演进。早期人们把像素值直接作为特征配合SVM、随机森林这类传统机器学习模型准确率已经不错。后来基于卷积神经网络的方案把识别率推到了很高的水平随便一个设计合理的网络都能在测试集上拿到超过99%的准确率。正因为数字识别足够成熟用它来练手才不会让算法创新干扰你学习框架本身。我在第一次接触PyTorch时也走了不少弯路。看教程觉得挺明白一写代码就报错尤其是张量形状、设备切换、梯度清零这些地方简直防不胜防。后来反复做这个数字识别项目才慢慢把许多概念真正固定在脑海里。所以说它不是“玩具项目”而是帮助你建立框架直觉的桥梁。1.2 做这个项目能帮你搞懂什么做完这个项目你会对PyTorch几个高频组件有实打实的理解而不是停留在“我好像见过”的层面。第一个是Dataset和DataLoader。你亲手创建了数据加载器就知道batch_size和shuffle到底在干什么。第二个是nn.Module。你从零定义一个模型类就会清楚__init__和forward的分工也不会再对着网络结构图发愁。第三个是训练循环。前向传播、计算损失、反向传播、更新参数这四步顺序如果有一步搞错模型就训练不起来做一遍胜过看十遍文字解释。除此之外你还会接触到张量的形状变化。比如一张图进入模型时是[1, 28, 28]一个批次进来是[64, 1, 28, 28]经过卷积和池化后变成[64, 16, 7, 7]最后展平成[64, 16*7*7]。这些维度为什么是这样的项目代码里的每个view、每个squeeze都会让你多一分手感。2. 数据集与预处理别小看这一步2.1 手写数字数据集的下载与加载PyTorch生态里最方便的地方就是torchvision已经内置了常用的视觉数据集和预处理工具。加载手写数字数据集只需要几行代码下载接口也封装好了省去了到处找数据包的麻烦。import torch import torchvision from torchvision import transforms # 定义一个预处理流水线注意顺序不能乱 transform transforms.Compose([ transforms.ToTensor(), # 将PIL图像转换为张量像素值从0~255映射到0~1 transforms.Normalize((0.5,), (0.5,)) # 再归一化到[-1, 1]公式为 (x - 0.5) / 0.5 ]) # 训练集下载并加载root指定缓存目录 train_dataset torchvision.datasets.MNIST( root./data, # 本地数据存放目录 trainTrue, # True表示训练集 transformtransform, # 每次取数据时自动应用预处理 downloadTrue # 如果本地没有自动下载 ) # 测试集trainFalse test_dataset torchvision.datasets.MNIST( root./data, trainFalse, transformtransform, downloadTrue )这段代码有几个细节值得你注意。root./data会在当前目录下创建一个data文件夹下载后的文件会缓存到这里下次再跑就不需要重复下载了。downloadTrue的底层逻辑是检查目标文件是否存在如果存在就跳过下载但如果你之前下载了不完整的文件有可能导致后续读取报错这时候把data目录下的压缩包删掉重新下载就行。transforms.Normalize((0.5,), (0.5,))这里的两个括号都是单元素因为灰度图只有一个通道。我之前见过有人把它写成Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))这是从三通道图像粘贴过来的代码跑起来立刻就会报通道数不匹配的错误。2.2 归一化和数据增强到底有没有必要很多初学者会疑惑ToTensor()已经把像素值从0到255缩放到0到1了为什么还要Normalize再来一次原因是神经网络训练时输入数据的范围直接影响梯度的稳定性。0到1之间虽然已经不错但如果把数据分布做到以0为中心比如-1到1对大多数激活函数来说更加友好梯度传播也更平稳。我实际对比过归一化和不归一化的效果。不归一化也能训练但损失曲线下降得慢同样的epoch数下准确率稍微低一些。如果你用的是比较深的网络或者BatchNorm之前的老式结构差异会更明显。所以别看这行代码不起眼它对训练速度和最终效果都有实打实的影响。至于数据增强数字识别任务里并不是必需品。训练集有六万张图像数量足够多而数字识别的关键是识别手写风格的变化这种变化可以通过卷积网络本身的平移等变性来应对一部分。我建议新手第一次先不加数据增强把基线跑通再尝试在transforms里加入随机旋转、随机平移。# 如果之后想尝试数据增强可以这样扩展 train_transform transforms.Compose([ transforms.RandomRotation(degrees10), # 随机旋转10度以内 transforms.RandomAffine(degrees0, translate(0.05, 0.05)), # 轻微平移 transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ])注意RandomRotation和RandomAffine必须在ToTensor之前因为它们操作的是PIL图像。如果顺序放反了程序会直接报错。3. 模型搭建从全连接到卷积3.1 设计一个能认出数字的神经网络模型结构的选择决定了你能学到什么程度的特征。对于28×28的小尺寸灰度图不用堆砌特别深的大模型一个两层的卷积网络已经足够展示卷积神经网络的核心思想。我设计的网络包含两个卷积层、两个最大池化层和一个全连接输出层。形象一点解释第一个卷积层负责识别图像的局部基础模式比如横线、竖线、弧线第二个卷积层把这些局部模式组合成更抽象的特征比如某个数字的上半部分、下半部分全连接层最终根据这些抽象特征判断属于哪一类数字。池化层做的是下采样把特征图尺寸缩小一半。它的作用不光是减少计算量更重要的是带来一定程度的平移不变性也就是说数字在图像里稍微偏移几个像素池化后提取的特征仍然比较稳定模型更能抓住本质。3.2 关键代码逐行注释下面给出完整的模型定义代码。我特意把注释写得比较啰嗦新手读起来会舒服很多。import torch.nn as nn import torch.nn.functional as F class DigitCNN(nn.Module): def __init__(self): super(DigitCNN, self).__init__() # 第一个卷积层输入1通道输出8通道卷积核3x3 # padding1表示在图像周围补一圈0这样卷积后大小不变 self.conv1 nn.Conv2d(in_channels1, out_channels8, kernel_size3, padding1) # 第二个卷积层输入8通道输出16通道卷积核3x3 # 通道数从8增加到16是为了组合更多特征 self.conv2 nn.Conv2d(in_channels8, out_channels16, kernel_size3, padding1) # 最大池化在2x2窗口内取最大值步长为2特征图尺寸减半 self.pool nn.MaxPool2d(kernel_size2, stride2) # 全连接层输入特征是16*7*7输出10个类别 self.fc1 nn.Linear(in_features16 * 7 * 7, out_features10) def forward(self, x): # 输入x形状: [batch_size, 1, 28, 28] x F.relu(self.conv1(x)) # 第一次卷积 ReLU: [b, 8, 28, 28] x self.pool(x) # 第一次池化: [b, 8, 14, 14] x F.relu(self.conv2(x)) # 第二次卷积 ReLU: [b, 16, 14, 14] x self.pool(x) # 第二次池化: [b, 16, 7, 7] x x.view(x.size(0), -1) # 展平: [b, 16*7*7] x self.fc1(x) # 全连接: [b, 10] return x这里有一个非常关键的细节为什么全连接层的输入维度是16*7*7因为28×28的图像经过两次2x2池化后尺寸先变成14×14再变成7×7输出通道数是16所以每个样本的特征就是一个16×7×7的立方体展平后就是16*7*7784。另外forward里没有出现softmax这是PyTorch的常见写法。模型只输出10个类别的原始分数logits损失函数CrossEntropyLoss会在内部自己计算softmax并处理数值稳定性。如果你在模型末尾又加了一层F.softmax反而会让损失计算变得不稳定属于画蛇添足。4. 训练流程损失函数、优化器与训练循环4.1 为什么选择交叉熵和Adam分类任务最自然的损失函数是交叉熵。它可以理解为衡量模型输出的概率分布与真实标签分布之间的差异差异越大损失值越高。在PyTorch中一句话就能初始化criterion nn.CrossEntropyLoss()这个损失函数有几个隐藏的细节。第一它接受两个输入模型的输出outputs形状是[batch_size, num_classes]真实标签labels形状是[batch_size]每个值是0到9的整数。它内部会自动把整数标签转成one-hot形式不需要你手动做编码。第二它把log_softmax和负对数似然合并在一起比单独调用softmax NLLLoss在数值上更稳定。优化器我推荐Adam对新手尤其友好。Adam的核心是自适应学习率它根据每个参数的历史梯度自动调整步长所以你不需要花太多时间在调学习率上设置lr0.001基本就能让模型正常收敛。相比之下传统SGD需要配合动量、合理的学习率调度新手容易一头雾水。4.2 训练循环的完整写法训练一个神经网络本质上就是反复执行四步操作前向传播计算预测、计算损失、反向传播算梯度、更新权重。下面是一个标准的训练循环。import torch.optim as optim # 初始化模型、损失函数、优化器 model DigitCNN() criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 创建数据加载器 train_loader torch.utils.data.DataLoader(train_dataset, batch_size64, shuffleTrue) # 训练轮数 num_epochs 5 for epoch in range(num_epochs): running_loss 0.0 # 遍历训练集中的所有批次 for images, labels in train_loader: # 第一步清空上一次迭代留下的梯度 optimizer.zero_grad() # 第二步前向传播得到预测分数 outputs model(images) # 第三步计算损失 loss criterion(outputs, labels) # 第四步反向传播计算每个参数的梯度 loss.backward() # 第五步根据梯度更新参数 optimizer.step() running_loss loss.item() # 打印本epoch的平均损失 avg_loss running_loss / len(train_loader) print(fEpoch {epoch1}/{num_epochs}, Loss: {avg_loss:.4f})这段代码里最容易被忽略的是optimizer.zero_grad()。PyTorch的设计是梯度默认累加如果你不清空梯度那么这一次的梯度会叠加到上一次的梯度上导致参数更新方向完全错乱。你可以试着把这一行去掉损失大概率会变成nan正是这个原因。还有一个细节loss.item()返回的是Python浮点数方便我们打印和统计而loss本身是一个张量它参与反向传播。如果你不小心在累加时写成了running_loss loss那么历史损失张量就会被长期持有占用的显存会不断增加epoch多了之后可能直接内存溢出。训练过程中你会看到损失从大概0.5慢慢下降到0.05附近。如果损失不降或者震荡剧烈通常要检查学习率是否过大或者网络结构是否有问题。4.3 保存和加载模型模型训练完之后一定要学会保存。PyTorch里推荐保存模型的state_dict也就是所有参数的权值字典而不是笼统地保存整个模型对象这样更灵活兼容性也更好。# 保存模型参数到文件 torch.save(model.state_dict(), digit_cnn.pth) # 之后需要加载时 model DigitCNN() # 先创建模型实例 model.load_state_dict(torch.load(digit_cnn.pth)) # 再加载参数注意load_state_dict之前必须有一个结构完全相同的模型实例。如果你修改了网络结构比如把第一层卷积的输出通道从8改成了16再加载旧的权重文件就会报错提示键名不匹配。这其实是个保护机制避免参数张量形状不符合预期。保存模型时也可以选择同时保存优化器的state_dict这样如果训练中断可以恢复学习率状态继续跑。5. 测试评估光跑完训练不够5.1 在测试集上计算准确率训练结束后我们必须用模型没见过的测试集来评估。很多新手在训练集上看了几个准确率就欢呼雀跃那并不能说明模型真的学到了泛化能力。# 创建测试集数据加载器shuffleFalse保持顺序 test_loader torch.utils.data.DataLoader(test_dataset, batch_size64, shuffleFalse) correct 0 total 0 # 测试阶段不需要计算梯度可以省显存和加速 with torch.no_grad(): for images, labels in test_loader: outputs model(images) # 取每个样本预测分数最高的类别索引 _, predicted torch.max(outputs, dim1) total labels.size(0) correct (predicted labels).sum().item() accuracy 100 * correct / total print(f测试集准确率: {accuracy:.2f}%)torch.max(outputs, dim1)会返回两个张量第一个是每一行的最大值第二个是最大值对应的索引。这里我们只关心索引所以用_接收第一部分。predicted labels会得到布尔张量.sum()统计True的数量.item()把它变成Python整数。这段代码跑完你正常情况下应该会看到准确率在98%到99%左右。如果明显低于95%建议回头检查数据预处理或训练循环是否写错。另外最好在测试前加上model.eval()虽然当前模型没有Dropout和BatchNorm层暂时不调用也不会出错但养成这个习惯很重要。同理训练阶段最好调用model.train()否则模型里某些层的行为会在训练和测试之间切换不干净。5.2 可视化预测结果直观发现问题光看一个数字不够直观。我每次测试都会随机打印一批图像把真实标签和预测标签放在一起这样哪类数字容易出错一目了然。import matplotlib.pyplot as plt # 从测试集中固定取一批 dataiter iter(test_loader) images, labels next(dataiter) # 模型预测 outputs model(images) _, predicted torch.max(outputs, dim1) # 画出前8张图 fig, axes plt.subplots(1, 8, figsize(12, 3)) for i in range(8): img images[i].squeeze() # 去掉单通道维度变成[28, 28] axes[i].imshow(img.numpy(), cmapgray) axes[i].set_title(f真:{labels[i].item()} 预测:{predicted[i].item()}) axes[i].axis(off) plt.show()图像张量在展示时需要先detach或直接在no_grad块内操作否则它带着梯度记录matplotlib会有警告甚至报错。上面代码使用了images[i].squeeze()因为单张灰度图的形状是[1, 28, 28]去掉维度1后才是[28, 28]。如果你想把多张图拼在一起可以回顾一下自己写的plt.subplots注意控制每个子图的尺寸避免重叠。5.3 用混淆矩阵观察分类错误准确率只能给你一个总体感受混淆矩阵能让你看到具体哪些类别容易被混淆。数字识别里3和5、4和9、7和1这类笔画相近的对是比较容易出错的。import numpy as np from sklearn.metrics import confusion_matrix # 如果不想安装sklearn也可以用简单的二维列表统计 confusion np.zeros((10, 10), dtypeint) with torch.no_grad(): for images, labels in test_loader: outputs model(images) _, predicted torch.max(outputs, dim1) for true_label, pred_label in zip(labels.numpy(), predicted.numpy()): confusion[true_label, pred_label] 1 print(混淆矩阵行真实标签列预测标签) print(confusion)如果你的混淆矩阵里对角线数值很大说明绝大多数预测正确如果某个格子比如confusion[4][9]特别大说明模型经常把手写4看成9。这时候你可以针对性收集那些失败样本看一下到底是图像模糊还是模型容量不够。我自己的经验是先用更大的训练集轮数调一调效果不明显再考虑改网络结构。6. 踩过的坑与调优经验6.1 数据加载时的维度问题新手最常见的错误就是维度对不上。torchvision返回的原始MNIST图像形状是[28, 28]经过ToTensor()后变成[1, 28, 28]DataLoader会再增加batch维度所以模型里的images实际上就是[batch_size, 1, 28, 28]。如果你在预处理时漏掉了ToTensor()传入模型的会是PIL图像直接导致后面所有计算爆炸。还有一种情况是我自己在做自定义数据集时踩过的图像通道数不是1而是3。有些手写数字图像是以RGB格式存储的虽然人眼看都一样但通道数变了模型第一层的in_channels1就会报错。解决办法要么把图像转成灰度要么把模型的输入通道数改成3。6.2 学习率设置不当导致不收敛学习率是训练里最“玄学”的参数但也是有规律可循的。Adam默认学习率0.001在小任务上通常表现很好但这不代表它永远不会出问题。如果你把学习率设成0.5损失值大概率会变成nan这是因为梯度过大导致参数更新太猛数值直接溢出。反过来学习率设成1e-6损失虽然不会报错但下降速度极慢训练了5个epoch准确率可能还是90%上下。我建议你在第一次跑的时候把lr固定在0.001等整个流程验证通过后可以尝试0.0001和0.01记录对应的损失曲线。如果损失震荡很大就调小如果下降太慢就调大。不要去盲猜损失曲线是最好的指南针。6.3 过拟合的判断与缓解训练集准确率接近100%测试集准确率却掉到95%这就是标准的过拟合。数字识别任务里过拟合往往来自于模型参数过多、训练轮数太长或者训练数据不足。虽然MNIST样本量已经很大但如果你的模型有几十万个参数依然存在记住训练样本的风险。缓解方法有三个层面。数据层面可以加旋转、平移、弹性变形等增强手段让模型看到更多变化模式模型层面可以减少卷积层通道数、加Dropout层训练层面可以减小训练轮数、加权重衰减。我自己试过把刚才那个模型的通道数从16加到64参数多了很多训练集准确率冲到99.8%但测试集准确率并没有明显提升还产生了波动。所以别一味追求大模型数字识别这个任务小而精的网络反而更实用。7. 后续还能怎么扩展7.1 用自己手写的图片测试模型模型训练完之后如果只用测试集验证总觉得不够过瘾。你完全可以自己拿纸笔写几个数字拍照或扫描下来用模型识别看看效果。from PIL import Image import torchvision.transforms as transforms # 读取手写数字图片并转为灰度 img Image.open(my_digit.png).convert(L) # L表示灰度模式 # 调整大小到28x28并应用和训练时一致的预处理 resize transforms.Resize((28, 28)) to_tensor transforms.ToTensor() normalize transforms.Normalize((0.5,), (0.5,)) img resize(img) img to_tensor(img) img normalize(img) # 增加batch维度因为模型需要四维输入 img img.unsqueeze(0) # 预测 model.eval() with torch.no_grad(): output model(img) _, predicted torch.max(output, dim1) print(f预测结果: {predicted.item()})这里有一个细节要注意你拍的图片可能自动带了黑边、白边或者其他噪声Resize直接压成28×28后数字会变形导致预测结果不准。更稳妥的做法是先裁剪出数字区域再填充到正方形最后缩放到28×28。我在实践时发现如果原图数字居中且笔画粗壮模型预测基本没问题如果原图数字偏小、有大量空白边缘模型容易出错。7.2 模型导出与部署思路如果你想把模型放到其他环境里用PyTorch提供了导出功能。最简单的是保存state_dict然后在新环境里重新定义同样的网络再加载权重。更高效的方式是导出为TorchScript这样就不需要依赖模型类定义了。# 导出为TorchScript脚本化模型 scripted_model torch.jit.script(model) scripted_model.save(digit_cnn.pt) # 加载时 loaded_model torch.jit.load(digit_cnn.pt)TorchScript的优势是模型变成自包含的文件部署时不需要再定义DigitCNN这个类适合放到服务端或移动端。不过导出前一定要确保模型已经调用model.eval()否则训练和测试模式相关的层会固化错误的行为。对初学者来说先掌握state_dict的保存加载就够了TorchScript留到有真实部署需求时再研究。7.3 从数字识别到更多视觉任务数字识别的框架几乎可以平移到所有单标签分类问题。换成彩色图像只需要把第一层卷积的in_channels从1改成3换成字母或交通标志只需要改数据集的加载方式和分类类别数想要识别图像里多个目标可以在此基础上引入目标检测网络比如基于区域或者基于锚框的方法但那些是另一个复杂度量级的项目了。如果你想把数字识别做成一个完整的小产品还可以给前面的图像加上一个简单的UI让用户用鼠标或者触摸板写数字然后实时调用模型识别。这个扩展既能练到PyTorch又能练到图形界面编程完成之后会特别有成就感。我个人做完这个小项目后最大的感受是先跑通再优化。很多人一开始就纠结到底用哪个网络结构、用不用迁移学习结果迟迟不敢动手。数字识别最宝贵的价值就是让你在半小时内经历从数据到模型的完整闭环。等你把每一行代码都敲熟了再回头去看更复杂的视觉项目你会发现它们的骨架其实和这个迷你项目非常相似。把这个流程刻在脑子里后续学任何新框架都会快很多。