PyTorch实战:手写数字识别从数据预处理到模型集成的完整流程
1. 从一次手写数字识别比赛说起手写数字识别这个题目放在今天的技术语境里多少有点老生常谈的味道。但有意思的是不管深度学习框架怎么迭代、模型结构怎么翻新这个任务始终是检验一个从业者基本功的试金石。我最近因为要带几个新人入门又重新把PyTorch做手写数字识别的完整流程走了一遍顺手把一些踩过的坑和优化细节整理出来。这篇文章面向的是刚接触PyTorch、想通过一个完整项目把深度学习流程跑通的朋友也适合那些做过类似任务但想看看别人怎么处理细节的老手。先说清楚这个任务本身给一张28x28像素的灰度图判断它是0到9中的哪个数字。训练集通常有几万张标注好的图片测试集是没标注的需要模型给出预测结果。听起来简单但里面涉及的东西一点都不少——数据加载、预处理、模型设计、训练循环、验证策略、超参数调整、结果提交每一个环节都有讲究。我见过太多人模型结构写得花里胡哨结果在数据标准化上栽跟头或者训练集准确率冲到99%但验证集死活上不去。这篇文章会按照实际操作的顺序来展开从环境准备到最终提交每一步都给出具体的代码和参数说明。更重要的是我会解释每个选择背后的原因——为什么用这个损失函数而不是那个为什么学习率要这样设为什么数据增强在这个任务上要谨慎使用。这些为什么才是真正能让你在下一个项目里举一反三的东西。2. 环境搭建与数据准备2.1 PyTorch环境配置的务实选择装PyTorch这件事说简单也简单说坑也多。我的建议是直接去官网的安装命令生成页面根据你的CUDA版本选择对应的命令。如果你没有GPU或者不想折腾驱动CPU版本完全够用——手写数字识别这个任务用CPU训练一个简单的CNN几分钟就能跑完一个epoch。# CPU版本安装示例 pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu # 如果有CUDA 11.8 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118这里有个细节值得说torchvision的版本要和torch匹配不然导入的时候会报一些莫名其妙的错误。我一般会先确认版本兼容性再开始写代码。另外如果你用conda环境建议单独建一个虚拟环境避免和系统里的其他包冲突。注意不要盲目追求最新版本。PyTorch的版本迭代很快但很多教程和预训练模型是基于特定版本写的。我一般会选择比最新版落后一两个小版本的稳定版踩坑概率低很多。2.2 数据加载与预处理的几个关键决策手写数字识别的数据集通常以CSV格式提供每行是一个样本第一列是标签后面784列是像素值。用pandas读进来之后需要做几件事import pandas as pd import numpy as np import torch from torch.utils.data import Dataset, DataLoader train_df pd.read_csv(train.csv) test_df pd.read_csv(test.csv) # 标签和像素分离 train_labels train_df[label].values train_images train_df.drop(label, axis1).values test_images test_df.values # 归一化到0-1之间 train_images train_images / 255.0 test_images test_images / 255.0 # 重塑为NCHW格式 train_images train_images.reshape(-1, 1, 28, 28) test_images test_images.reshape(-1, 1, 28, 28)这里有几个决策点需要解释。第一为什么要除以255因为原始像素值是0到255的整数直接输入网络会导致数值过大梯度更新不稳定。归一化到0到1之间是最基础的操作但很多人会忘记。第二为什么要reshape成NCHWPyTorch的卷积层要求输入是(batch, channel, height, width)的格式而CSV里每行是扁平的784个值必须手动重塑。还有一个容易被忽略的点数据类型。train_images经过除法之后是float64但PyTorch默认使用float32。如果不转换后续计算会报类型不匹配的错误。我一般会在创建Dataset的时候统一转成float32。class DigitDataset(Dataset): def __init__(self, images, labelsNone, transformNone): self.images images.astype(np.float32) self.labels labels self.transform transform def __len__(self): return len(self.images) def __getitem__(self, idx): image self.images[idx] if self.transform: image self.transform(image) if self.labels is not None: return image, self.labels[idx] return image2.3 验证集划分的策略很多人做比赛或者练习的时候喜欢把全部训练数据拿来训练然后靠提交结果看分数。这种做法在真正的项目里是不可取的因为你无法判断模型是真的学到了东西还是过拟合了。我的习惯是从训练集里划出10%到20%作为验证集用来监控训练过程。from sklearn.model_selection import train_test_split X_train, X_val, y_train, y_val train_test_split( train_images, train_labels, test_size0.15, random_state42, stratifytrain_labels )stratify参数很重要它保证训练集和验证集里各个数字的比例一致。手写数字识别虽然各类样本相对均衡但养成分层抽样的习惯没坏处。random_state固定住保证每次运行划分结果一样方便对比不同模型的性能。3. 模型设计与训练循环3.1 为什么选CNN而不是全连接网络手写数字识别最基础的方案是用全连接网络把784个像素拉平接几个隐藏层最后输出10个类别的概率。这个方案能跑但效果一般。原因在于全连接网络丢失了图像的空间结构信息——它不知道相邻像素之间的关系也不知道某个像素在图像的哪个位置。卷积神经网络通过局部连接和权值共享解决了这个问题。一个3x3的卷积核在图像上滑动提取的是局部特征比如边缘、角点。多层卷积堆叠之后网络能学到越来越抽象的特征。对于28x28的小图一个简单的CNN结构就足够了。我常用的一个基准结构是这样的import torch.nn as nn import torch.nn.functional as F class DigitCNN(nn.Module): def __init__(self): super(DigitCNN, self).__init__() self.conv1 nn.Conv2d(1, 32, kernel_size3, padding1) self.bn1 nn.BatchNorm2d(32) self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) self.bn2 nn.BatchNorm2d(64) self.pool nn.MaxPool2d(2, 2) self.dropout1 nn.Dropout(0.25) self.fc1 nn.Linear(64 * 7 * 7, 128) self.dropout2 nn.Dropout(0.5) self.fc2 nn.Linear(128, 10) def forward(self, x): x self.pool(F.relu(self.bn1(self.conv1(x)))) x self.pool(F.relu(self.bn2(self.conv2(x)))) x x.view(-1, 64 * 7 * 7) x self.dropout1(x) x F.relu(self.fc1(x)) x self.dropout2(x) x self.fc2(x) return x这个结构里每个选择都有理由。卷积核用3x3是最常见的因为它在参数量和感受野之间取得了很好的平衡。padding1保证卷积后特征图尺寸不变方便后续计算。BatchNorm放在卷积之后、激活之前能加速收敛并起到一定的正则化作用。MaxPooling用2x2每次把特征图尺寸减半。Dropout放在全连接层之前防止过拟合。3.2 损失函数与优化器的选择逻辑分类任务用交叉熵损失这个没什么争议。PyTorch的CrossEntropyLoss内部已经包含了Softmax操作所以模型最后一层不需要加Softmax。criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3)优化器我选Adam而不是SGD原因是Adam对学习率不那么敏感默认的1e-3在大多数情况下都能工作得不错。SGD虽然在某些任务上能收敛到更好的解但需要仔细调学习率和动量对新手不友好。如果你追求极致性能可以试试SGD加余弦退火但那是后话。学习率调度器我一般会加一个ReduceLROnPlateau当验证集损失不再下降时自动降低学习率scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience3, verboseTrue )patience3表示连续3个epoch验证损失没改善就降低学习率。factor0.5表示每次降低一半。这个策略比固定学习率稳得多。3.3 训练循环的完整实现与监控训练循环的代码看起来都差不多但细节决定成败。我习惯在每个epoch结束后打印训练损失、训练准确率、验证损失、验证准确率并保存验证准确率最高的模型。def train_epoch(model, loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for images, labels in loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted outputs.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() return running_loss / total, correct / total def validate(model, loader, criterion, device): model.eval() running_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) running_loss loss.item() * images.size(0) _, predicted outputs.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() return running_loss / total, correct / total这里有个容易犯的错误忘记在验证时调用model.eval()和torch.no_grad()。前者影响Dropout和BatchNorm的行为后者影响内存占用和计算速度。不加no_grad的话验证过程会构建计算图显存占用会飙升。实操心得我习惯在训练循环里加一个早停机制。如果验证损失连续10个epoch没有改善就停止训练。这能省下大量时间也能防止过拟合。4. 性能优化与结果提交4.1 数据增强的谨慎使用数据增强是提升模型泛化能力的常用手段但在手写数字识别这个任务上要小心。旋转、缩放、平移这些操作如果幅度过大可能会把6变成9把1变成7。我一般只用小幅度的随机旋转±10度和随机平移±10%而且会先可视化增强后的图片确认标签仍然正确。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomRotation(10), transforms.RandomAffine(degrees0, translate(0.1, 0.1)), transforms.ToTensor(), ])注意ToTensor()会把像素值从0-1缩放到0-1如果输入是PIL图像或者保持原样如果输入是numpy数组。如果你已经在Dataset里做了归一化这里就不需要重复操作。4.2 学习率与批大小的调参经验批大小我一般设64或128。太小的话训练不稳定太大的话泛化能力可能下降。学习率用1e-3起步如果训练损失震荡就降到1e-4。如果训练损失下降太慢就升到3e-3试试。有一个经验性的规律批大小翻倍学习率也可以翻倍。但这个规律不是绝对的还是要看具体的训练曲线。我通常会跑几个小实验每个实验只训练5个epoch看看损失下降的趋势再决定最终的参数。4.3 模型集成与结果提交单模型的准确率通常能到99%左右如果想再往上提一点可以用模型集成。最简单的做法是训练多个结构相同但初始化不同的模型然后把它们的预测概率平均。def ensemble_predict(models, test_loader, device): all_probs [] for model in models: model.eval() probs [] with torch.no_grad(): for images in test_loader: images images.to(device) outputs model(images) probs.append(F.softmax(outputs, dim1).cpu().numpy()) all_probs.append(np.concatenate(probs)) avg_probs np.mean(all_probs, axis0) return np.argmax(avg_probs, axis1)集成3到5个模型通常能把准确率提升0.1到0.3个百分点。虽然提升不大但在比赛里往往就是这点差距决定排名。提交格式一般是CSV两列ImageId和Label。ImageId从1开始Label是预测的数字。submission pd.DataFrame({ ImageId: range(1, len(predictions) 1), Label: predictions }) submission.to_csv(submission.csv, indexFalse)5. 常见问题与排查技巧实录5.1 训练损失不下降的排查思路这是最常见的问题。我一般按以下顺序排查排查项检查方法常见原因数据标签打印前几个样本的标签和图片标签和图片错位数据归一化检查像素值范围忘记除以255学习率尝试1e-4到1e-2太大导致震荡太小导致不收敛模型输出打印logits的均值和方差初始化有问题损失函数确认用的是CrossEntropyLoss误用MSELoss我遇到过一次训练损失死活不降最后发现是DataLoader的shuffle设成了False而且数据是按标签排序的模型每次只看到同一个数字根本学不到东西。5.2 验证集准确率远低于训练集的应对这是典型的过拟合。解决方法按优先级排序增加Dropout比例从0.25提到0.5加L2正则化在优化器里设weight_decay1e-4减少模型参数量比如把全连接层的128改成64增加训练数据或者用数据增强早停在验证损失开始上升时就停止训练注意不要一上来就加各种正则化。先确认训练集和验证集的分布是否一致有时候是数据划分的问题不是模型的问题。5.3 GPU内存不足的处理如果你用GPU训练可能会遇到CUDA out of memory。解决方法减小批大小从128降到64或32用torch.cuda.empty_cache()清理缓存在验证时确保用了torch.no_grad()如果模型太大考虑用梯度累积模拟大批大小梯度累积的写法accumulation_steps 4 for i, (images, labels) in enumerate(loader): outputs model(images) loss criterion(outputs, labels) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这样等效于批大小扩大了4倍但显存占用不变。5.4 预测结果全为同一个数字的诡异情况这个问题我遇到过两次一次是因为测试集的数据没有做和训练集一样的归一化另一次是因为模型最后一层忘了加Softmax虽然CrossEntropyLoss内部有但如果你手动算概率就会出错。排查的时候先打印测试集前几个样本的预测概率分布如果某个类别的概率接近1其他接近0那说明模型可能过拟合了如果所有类别的概率都差不多那说明模型没学到东西。6. 一些不在文档里的实操体会跑完整个流程之后我有几个体会想分享。第一不要迷信复杂的模型结构。在这个任务上一个三层卷积加两层全连接的简单网络和那些十几层的ResNet准确率差距可能不到0.5%。但训练时间和调试难度差了好几倍。第二数据质量比模型结构重要。我试过用同样的模型在清理过的数据上训练比在原始数据上训练验证准确率高出一个百分点。第三记录实验很重要。我习惯用表格记录每次实验的超参数和结果这样能快速定位哪些改动有效哪些是白费功夫。最后分享一个小技巧如果你在本地训练可以把模型保存下来下次直接加载继续训练。PyTorch的state_dict机制很方便torch.save(model.state_dict(), model.pth) model.load_state_dict(torch.load(model.pth))这样即使训练中断也不用从头再来。对于手写数字识别这种小任务可能感觉不到优势但养成这个习惯以后做大规模训练的时候会感谢自己。