最近有朋友问我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', # 本地数据存放目录 train=True, # True表示训练集 transform=transform, # 每次取数据时自动应用预处理 download=True # 如果本地没有，自动下载 ) # 测试集：train=False test_dataset = torchvision.datasets.MNIST( root='./data', train=False, transform=transform, download=True )这段代码有几个细节值得你注意。root='./data'会在当前目录下创建一个data文件夹，下载后的文件会缓存到这里，下次再跑就不需要重复下载了。download=True的底层逻辑是检查目标文件是否存在，如果存在就跳过下载，但如果你之前下载了不完整的文件，有可能导致后续读取报错，这时候把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(degrees=10), # 随机旋转10度以内 transforms.RandomAffine(degrees=0, 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 # padding=1表示在图像周围补一圈0，这样卷积后大小不变 self.conv1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, padding=1) # 第二个卷积层：输入8通道，输出16通道，卷积核3x3 # 通道数从8增加到16，是为了组合更多特征 self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, padding=1) # 最大池化：在2x2窗口内取最大值，步长为2，特征图尺寸减半 self.pool = nn.MaxPool2d(kernel_size=2, stride=2) # 全连接层：输入特征是16*7*7，输出10个类别 self.fc1 = nn.Linear(in_features=16 * 7 * 7, out_features=10) 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*7=784。
另外，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的核心是自适应学习率，它根据每个参数的历史梯度自动调整步长，所以你不需要花太多时间在调学习率上，设置lr=0.001基本就能让模型正常收敛。相比之下，传统SGD需要配合动量、合理的学习率调度，新手容易一头雾水。
4.2 训练循环的完整写法
训练一个神经网络本质上就是反复执行四步操作：前向传播计算预测、计算损失、反向传播算梯度、更新权重。下面是一个标准的训练循环。
import torch.optim as optim # 初始化模型、损失函数、优化器 model = DigitCNN() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # 创建数据加载器 train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True) # 训练轮数 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(f'Epoch {epoch+1}/{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 在测试集上计算准确率
训练结束后，我们必须用模型没见过的测试集来评估。很多新手在训练集上看了几个准确率就欢呼雀跃，那并不能说明模型真的学到了泛化能力。
# 创建测试集数据加载器，shuffle=False保持顺序 test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False) correct = 0 total = 0 # 测试阶段不需要计算梯度，可以省显存和加速 with torch.no_grad(): for images, labels in test_loader: outputs = model(images) # 取每个样本预测分数最高的类别索引 _, predicted = torch.max(outputs, dim=1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = 100 * correct / total print(f'测试集准确率: {accuracy:.2f}%')torch.max(outputs, dim=1)会返回两个张量：第一个是每一行的最大值，第二个是最大值对应的索引。这里我们只关心索引，所以用_接收第一部分。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, dim=1) # 画出前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(), cmap='gray') 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), dtype=int) with torch.no_grad(): for images, labels in test_loader: outputs = model(images) _, predicted = torch.max(outputs, dim=1) 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_channels=1就会报错。解决办法要么把图像转成灰度，要么把模型的输入通道数改成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, dim=1) 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，又能练到图形界面编程，完成之后会特别有成就感。
我个人做完这个小项目后最大的感受是：先跑通，再优化。很多人一开始就纠结到底用哪个网络结构、用不用迁移学习，结果迟迟不敢动手。数字识别最宝贵的价值就是让你在半小时内经历从数据到模型的完整闭环。等你把每一行代码都敲熟了，再回头去看更复杂的视觉项目，你会发现它们的骨架其实和这个迷你项目非常相似。把这个流程刻在脑子里，后续学任何新框架都会快很多。
