我用PyTorch和CNN把MNIST手写数字识别从头到尾做了一遍。这篇博文从数据管道讲到网络结构、训练循环和调参心得,包括不少我第一次写这段代码时踩进去的坑,适合刚接触深度学习的读者照着复现,也能给已经在跑模型的人一些排查思路。
1. 项目拆解:MNIST识别为什么选CNN而不是全连接
1.1 数据形态与任务本质:28x28的灰度图意味着什么
MNIST是手写数字识别里最经典的入门数据集。训练集有60000张图,测试集有10000张图,每张都是28x28像素的灰度图,像素值范围在0到255之间,标签是0到9这十个数字。这个数据集规模不大、类别固定、背景干净,所以特别适合用来验证一个模型结构或者训练流程是不是靠谱,也是很多人接触深度学习时跑通的第一个完整项目。
这里有个容易被忽略的关键点:如果把每张图当成一个向量输入全连接网络,那就是784个特征,按说特征也不多,但效果就是不如CNN好。原因在于图像数据的特性——数字是由局部的笔画、边缘、拐角组合出来的,这些局部结构在二维平面上有明确的空间关系。全连接网络把784个像素抻平了,空间邻接关系被强行切断,模型只能学到"像素点之间的统计相关性",学不到"笔画在图上怎么排布"。比如数字7和数字1,如果只看像素分布,某些写法差异极小,但它们在图上占据的位置和形状完全不同。CNN天然用二维卷积核扫描图像,每次只关注一个小邻域,局部特征提取完再逐层抽象,这和人类看数字的方式更接近。
MNIST虽然简单,但它是理解CNN的绝佳载体。这个数据集上,一个结构合理的CNN在测试集上跑出99%以上的准确率是稳定能做到的,训练时间在普通CPU上也就几分钟到十几分钟。也就是说,我们不需要担心算力不够,可以把全部注意力放在理解卷积、池化、训练流程这些核心环节上。我觉得这是MNIST最有价值的地方。
1.2 CNN的三个杀手锏:局部连接、权值共享、平移不变性
CNN相对全连接网络有三大核心机制,都在MNIST这个简单任务上体现得淋漓尽致。
第一个是局部连接。全连接层里每个输出神经元要跟上一层所有神经元相连,一个28x28的图像展开成784维向量,第一层如果有256个神经元,那就有784x256约20万个参数。CNN不一样,卷积核比如3x3,每个输出像素只跟上一层对应的3x3邻域相连,参数数量和图像尺寸无关,只跟卷积核大小和通道数有关。这相当于用一个"局部视野"扫描全图,而不是让每个神经元都看全图。
第二个是权值共享。同一个卷积核在整张图上滑动,用的都是同一组权重。假设输入是28x28单通道,用32个3x3卷积核,参数量就是32x3x3+32(偏置),才320个。相比之下,同样输出32个特征图的全连接层需要784x32约2.5万个参数。权值共享的本质是假设"某种局部特征(比如横线)在图像任何位置出现都有相同的含义",这非常契合图像数据的特点。
第三个是平移不变性。数字稍微往左挪几个像素、往右挪几个像素,它还是同一个数字。卷积操作搭配池化,尤其是最大池化,让模型对小幅平移、轻微形变有很强的容忍度。池化层取窗口内的最大值,相当于保留"这个区域里有某个特征",而不关心它精确落在哪个像素位置。
这三个机制让CNN在图像任务上的参数效率远高于全连接网络。MNIST上用全连接网络也能跑到97%左右,但要达到99%,全连接需要非常宽的网络和大量正则化,CNN只需要一个很小的网络结构就能轻松超过它。我们自己做项目的时候,不需要迷信"网络越大越好",理解了这个对比,就知道结构设计对参数量和泛化能力的影响有多大。
1.3 整体方案选型:从框架到网络结构的核心取舍
做这个项目之前,先把方案的关键决策理清楚。
框架选型上,PyTorch是现在研究和教学里最常用的选择之一。它的动态图机制让网络结构可以随时用Python的if、for控制,调试的时候非常直观;配合张量维度检查和训练过程可视化,新手定位问题会舒服很多。TensorFlow也很好,但PyTorch的调试体验和生态适配度更适合做这类快速迭代的项目。
网络结构上,我采用"卷积层+池化层交替、最后接全连接层"的经典模式。第一层卷积把单通道灰度图映射到多通道特征图,第二层卷积进一步抽象,最后把二维特征图展平成一维向量,交给全连接层做分类。隐藏层激活函数用ReLU,输出层配合交叉熵损失,因为PyTorch的CrossEntropyLoss里面内置了Softmax,不需要在模型最后一层手动加Softmax。
超参数初始化方面,我选择批量大小64、学习率0.001、训练10个epoch。这三个数值是大量MNIST实践里被验证过的稳妥组合。学习率0.001配合Adam优化器,收敛平稳;批量大小64在性能和显存占用之间比较均衡;10个epoch足够让模型收敛到99%附近,又不至于过度训练浪费时间。后面我会详细解释这些选择的依据,以及怎么根据自己的硬件情况动态调整。
2. 环境准备与数据管道搭建
2.1 版本选型与安装踩坑
环境这块是很多初学者第一个卡住的地方。我的建议是:装PyTorch的时候不要直接pip install torch拉最新版,先看一眼CUDA版本。我最早就是这么干的,结果装了CPU版,后面用GPU训练还得重装一遍。正确的做法是去PyTorch官网的安装向导页面,选好自己的操作系统、包管理方式、CUDA版本,它生成的命令里会带合适的下载源。如果暂时不打算用GPU,CPU版也完全能跑MNIST,就是每个epoch慢个一两分钟,完全在可接受范围内。
还需要装torchvision,因为MNIST数据集和常用的图像预处理工具都在这个包里。版本上只要跟PyTorch主版本匹配就行,通常写pip install torch torchvision一把装齐。装完之后在Python里跑一句import torch验证,如果报缺少DLL或者版本不匹配之类的错,优先排查是不是装了多套Python环境。我习惯用虚拟环境隔离项目,这样不管怎么折腾都不会污染系统环境。
数据这块最典型的坑是下载速度。MNIST数据集本身很小,但有些地区访问数据集服务器的速度很慢。torchvision.datasets.MNIST默认会在首次访问时自动下载到某个目录,如果下载失败,可以手动从镜像源把四个文件(train-images、train-labels、test-images、test-labels)下载好,解压后放到./data/MNIST/raw/目录下,代码里设置download=True就不会再去下载了。提前准备好数据文件,能省掉不少等待时间。
2.2 数据加载与预处理:归一化为什么用0.1307和0.3081
数据预处理是深度学习流程里看似简单但极其重要的一环。MNIST原始像素值范围在0到255,如果直接喂给网络,数值太大,会让梯度更新不稳定,收敛变慢。常见的做法是做标准化,把像素值缩放到均值为0、方差为1的分布附近。
PyTorch里用torchvision.transforms.ToTensor()先把PIL图像转成张量,像素值自动从0~255缩放到0~1。然后接transforms.Normalize((0.1307,), (0.3081,))。第一个参数0.1307是MNIST数据集的像素均值,第二个0.3081是标准差。这两个数值不是随便拍的,是统计整个训练集得到的。标准化之后,像素值大致分布在-1到1之间,这会帮助模型更快收敛。
很多人在这一步会犯一个错误:用了归一化却不知道这两个数字是干什么的,或者自己随便填了个均值标准差。MNIST官方数据集在torchvision里就有默认的统计值,直接抄过来用完全没问题。但如果自己处理其他数据集,务必先算一下均值和标准差,公式是(x - mean) / std逐通道计算。再说一句,归一化里用到的均值和标准差应该来自训练集,测试集同样用训练集的统计结果,不能在测试集上重新算,否则就是数据泄漏。
预处理这块我完整贴一下代码,后面所有步骤都基于这个数据管道:
import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST( root='./data', train=True, download=True, transform=transform ) test_dataset = datasets.MNIST( root='./data', train=False, download=True, transform=transform )Compose把多个预处理操作串起来,执行顺序就是列表里写的顺序。先转张量,再归一化,这个顺序不能反过来,因为Normalize只能处理张量,不能处理PIL图像。
2.3 DataLoader与批训练机制:为什么一次喂一批而不是一张一张喂
数据加载的下一步是用DataLoader把数据集包装成可迭代对象。我一般这样配置:
train_loader = DataLoader( train_dataset, batch_size=64, shuffle=True ) test_loader = DataLoader( test_dataset, batch_size=64, shuffle=False )这里有几个值得展开说的细节。
shuffle=True只在训练集上设置。每个epoch开始前把数据打乱,让模型每次看到的样本顺序都不同,避免因为原始数据的排序规律导致训练偏差。测试集不需要打乱,因为测试只是评估,不打乱还可以方便后续可视化每个batch里的样本和对应的预测结果。
batch_size=64表示一次前向传播同时处理64张图。为什么不一张一张喂呢?因为GPU或CPU的并行计算特性决定了批量处理效率更高,矩阵运算可以一次算64个样本,充分利用硬件计算单元。那batch_size是不是越大越好呢?不是。太大的batch会让模型在梯度下降时更稳定但更容易陷入尖锐极小值,泛化能力反而可能变差;太小的batch梯度噪声大,收敛不稳定。64到128之间是MNIST上的常见选择。
迭代次数有个小知识要理清。很多人第一次写循环会把epoch和iteration搞混。一个epoch指的是完整遍历一遍训练集。如果训练集有60000张图,batch_size是64,那么一个epoch里有60000/64约938个iteration(最后一批不足64会取较少样本)。跑10个epoch,模型一共看到60000x10次样本更新。这个关系后面写训练循环的时候会直接用到,理解了就不会在处理"跑了几轮"的问题上犯迷糊。
3. CNN架构设计与维度推导
3.1 网络结构整体设计:两层卷积加三层全连接
这个项目我用一个精简的CNN结构,它不是随便搭的,而是借鉴了经典卷积网络的设计思路,又为MNIST做了尺寸适配。
import torch.nn as nn class CNN(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(64 * 7 * 7, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.pool(torch.relu(self.conv1(x))) x = self.pool(torch.relu(self.conv2(x))) x = x.view(x.size(0), -1) x = torch.relu(self.fc1(x)) x = self.fc2(x) return x结构拆开看是这样:
- conv1:输入1通道,输出32通道,3x3卷积核,padding=1保持尺寸不变。
- pool:2x2最大池化,步长2,图像尺寸减半。
- conv2:输入32通道,输出64通道,同样3x3卷积核和padding。
- pool:再次减半,此时特征图从28x28变成7x7。
- fc1:把64个7x7特征图展平,得到64x7x7=3136个特征,全连接到128维。
- fc2:128维映射到10个类别。
为什么用32和64这两个通道数?这是参数量和表达能力的折中。通道数太少,模型表达能力不足;通道数翻倍是CNN常见的做法,因为随着空间尺寸变小,可以增加通道数来弥补空间信息的损失,让网络保持足够的信息容量。MNIST类别少、图像简单,32和64已经足够,再往上加到128甚至256也不是不行,但收益有限,训练时间反而变长。
最后输出层为什么没有显式加Softmax?这里很多人会疑惑。因为PyTorch的nn.CrossEntropyLoss在计算损失时,内部会对模型的原始输出(logits)做Softmax归一化,再计算交叉熵。如果我在模型里先做了Softmax,再去算交叉熵,反而会造成数值不稳定且结果错误。新手最容易在这里被绕晕,记住一点:训练时模型输出logits,损失函数内部处理Softmax;只有推理时需要概率输出,才手动加Softmax。
3.2 卷积参数量计算:为什么卷积层参数这么少
计算一下这个网络的参数量,你会对CNN的参数效率有更直观的认识。
conv1层的参数量是:输出通道32乘以输入通道1乘以卷积核尺寸3x3,再加上32个偏置。算下来32x1x9+32=320个参数。conv2层参数量是:64x32x9+64=18496个参数。两层卷积加起来大约1.88万个参数。
全连接层那边,fc1输入是64x7x7=3136,输出128,参数量3136x128+128约40.2万;fc2是128x10+10=1290。整个网络总参数量大约42万。
这42万参数里,全连接层贡献了九成以上,卷积层虽然做了主要特征提取工作,参数占比反而不高。这也是现代网络设计里经常用全局平均池化替代全连接层的原因之一,能大幅减少尾部参数。但对于MNIST这种小任务,尾部全连接层完全够用,不用刻意追求极致的参数效率。
对比一下全连接网络:一个简单的两层全连接网络,输入784维,隐藏层256,输出10,参数量就是784x256+256约20.1万,再加上256x10+10约2570,总共20万左右。看起来比42万少,但效果远不如CNN。CNN用更合理的结构大幅提升每参数的有效性。理解了参数量构成,你在设计自己的网络时就会清楚瓶颈在哪、该从哪里优化。
3.3 维度推导全流程:一张图从28x28到10个输出
维度不匹配是PyTorch新手最常见的报错来源。与其等报错出来再猜,不如在写代码的时候就把每一步的维度变化算清楚。我们跟一张28x28的图完整走一遍。
输入形状是[1, 28, 28],1表示通道数(灰度图),28x28是宽高。进入网络后,一次处理一个batch,所以实际形状是[64, 1, 28, 28],64就是batch_size。
第一层conv1,输入[64, 1, 28, 28],输出[64, 32, 28, 28]。因为kernel_size=3、padding=1、stride=1,输出尺寸公式是(W - F + 2P) / S + 1,代入就是(28 - 3 + 2) / 1 + 1 = 28,空间尺寸不变,通道从1变成32。
然后ReLU和最大池化,池化核2x2、步长2,输出尺寸是28 / 2 = 14,得到[64, 32, 14, 14]。
第二层conv2,输入[64, 32, 14, 14],输出[64, 64, 14, 14]。通道从32变成64,空间尺寸不变。
再池化一次,14除以2等于7,得到[64, 64, 7, 7]。
这里就要解释fc1 = nn.Linear(64 * 7 * 7, 128)里的6477是怎么来的。展平操作x.view(x.size(0), -1)把每个样本从[64, 7, 7]变成一维向量64*7*7=3136,批量维度保留在开头。全连接层输入就是3136。
fc1输出[64, 128],fc2输出[64, 10]。到这里,64个样本每个对应10个类别的得分,维度就全部通了。
这个推导过程每一步都可以用print(x.shape)验证。我第一次写的时候就是因为对池化之后的尺寸变化没算清楚,fc1输入维度写错,报了一堆维度错误。后来养成了一个习惯,设计网络的时候先在纸上推一遍维度,再写代码,基本就告别这类报错了。
4. 训练配置与核心循环实现
4.1 损失函数与优化器选择:CrossEntropyLoss和Adam的配合逻辑
训练一个分类模型,损失函数和优化器的选择直接决定能否收敛以及收敛速度。
损失函数我选nn.CrossEntropyLoss()。它做的事情是:先对logits做Softmax转成概率分布,然后计算预测分布和真实标签的交叉熵。交叉熵的特点是,当模型对正确类别的概率接近1时损失趋近于0,当模型混淆类别时损失会给出较大的惩罚信号。相比均方误差,交叉熵在分类任务上收敛更快,因为它的梯度不会因为Softmax的饱和效应而消失。具体来说,Softmax后的概率如果非常接近0或1,MSE的梯度会变得很小,而交叉熵的梯度是softmax输出 - one-hot真实标签,即使输出饱和,梯度也不会趋近于零。
优化器我选optim.Adam,学习率1e-3。Adam是自适应矩估计优化器,它会为每个参数单独维护学习率,基于梯度的一阶矩和二阶矩估计动态调整。简单说,它对学习率不那么敏感,即使你选的学习率不太合适,它也大概率能收敛。我第一次训练神经网络的时候直接用了SGD,学习率设了0.1,结果训练时损失一直在震荡,后来换成Adam加1e-3,几步就稳定下降了。对于这类入门项目,Adam是最稳妥的选择。
SGD也不是不能用。SGD配合动量在训练后期往往能取得比Adam更好的泛化效果,但它的学习率和动量参数需要更仔细的调试。如果你想深入研究优化器对模型性能的影响,可以在模型收敛后用SGD微调几十个epoch,或者把SGD的lr设得很小再全局跑一遍。但对项目本身的目的是"完整跑通一个CNN流程"这个阶段来说,Adam的工程稳定性是最重要的。
import torch.optim as optim criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001)model.parameters()会把模型里所有可训练的参数交给优化器管理,包括卷积核权重、偏置、全连接层权重和偏置。这里不需要手动区分哪些参数参与训练,默认都要更新。
4.2 训练循环与验证流程:一个epoch的三个关键阶段
训练循环是项目的心脏。我贴出完整代码,然后分阶段解释每一行的作用。
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = CNN().to(device) def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() total_loss = 0 correct = 0 total = 0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() avg_loss = total_loss / len(train_loader) accuracy = 100.0 * correct / total print(f'Train Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%') def evaluate(model, test_loader, device): model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in test_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = 100.0 * correct / total print(f'Test Accuracy: {accuracy:.2f}%') return accuracy for epoch in range(10): print(f'Epoch {epoch+1}/10') train_one_epoch(model, train_loader, criterion, optimizer, device) evaluate(model, test_loader, device)有几个容易被忽略的习惯非常关键。
第一,model.train()和model.eval()不是摆设。这两个模式会影响某些层的运行逻辑,比如Dropout在训练时随机丢弃神经元、在评估时不做丢弃;BatchNorm在训练时用当前batch的均值方差归一化,评估时用累积的全局统计量。虽然当前模型里没有这些层,但养成交替调用这两个方法的习惯,后面模型复杂了才不会踩坑。
第二,optimizer.zero_grad()必须在每次反向传播之前清空梯度。PyTorch的梯度是累积的,不清空的话,下一个batch的梯度会叠加到上一个batch上,导致更新方向完全错误。这个顺序必须是:清空梯度、前向计算、计算损失、反向传播、更新参数。
第三,torch.no_grad()在评估阶段关闭梯度计算。评估过程不需要反向传播,如果开着梯度,会额外消耗显存和计算资源。虽然MNIST很小,影响不显著,但这是工程上必备的效率意识。
第四,torch.max(outputs, 1)这个函数取每个样本在10个类别得分中的最大值和索引,索引就是预测类别。(predicted == labels).sum().item()把预测正确的数量统计出来。这里要注意像素值计算完后要.item()取出Python数值,直接用张量做累加会导致计算图被保留,浪费内存。
4.3 训练结果观察:从损失曲线判断收敛状态
训练10个epoch之后,典型的输出是训练准确率从90%左右一路升到99%以上,测试准确率最终稳定在99%以上。我自己跑的时候,第一个epoch结束训练准确率大概93%,测试准确率约95%;第5个epoch训练接近99%;第10个epoch测试准确率99.2%左右。
如果你看到的曲线和这个差距很大,可以从这几个角度排查。损失一直在降但准确率不涨,先看数据集标签对不对;训练准确率到了99%但测试准确率只有97%,这是典型的过拟合,说明模型在训练集上记住了样本,泛化能力不够;损失一开始就是nan,通常是学习率过大导致梯度爆炸,调低学习率就可以解决。
每轮epoch打印一次训练损失和测试准确率,是监控训练状态的最低配置。如果希望看得更细,可以把每个iteration的损失记录下来,画成曲线。MNIST上跑10个epoch,每个epoch940个batch左右,所以每50个iteration打印一次损失就能看出下降趋势,又不会刷屏。我后来做复杂项目时,都会用TensorBoard记录损失曲线和权重分布,但在这个入门项目里,简单的print输出已经够用。
5. 常见问题与调参实战
5.1 维度不匹配:PyTorch最常见的报错之一
我见过太多人在MNIST这个项目上报RuntimeError: mat1 and mat2 shapes cannot be multiplied这类错误。根本原因就一个:全连接层输入维度和展平后的特征数对不上。
比如你把fc1定义成nn.Linear(64 * 7 * 7, 128),但实际池化后的特征图是[64, 32, 14, 14](如果少了一层池化),那view展平后是6272维,和64*7*7=3136对不上,立刻报错。遇到这种错误,最快的排查方法是在模型forward里加print(x.shape),从输入开始每一步把形状打出来,找到第一个跟预期不符的位置,就知道是哪一层设计出了问题。
还有一种情况是view里用了x.view(-1, 3136)。这种写法把批量维也揉进去了,在inference阶段如果一次只喂一张图,形状变成[1, 3136]没问题,但训练时batch_size=64,就会变成[64, 3136]和[16384, ...]这种莫名其妙的形状。正确写法永远是x.view(x.size(0), -1),保留批量维度,自动展平剩余部分。这个错误特别隐蔽,因为在某些条件下能跑通,换了batch_size就崩。
5.2 训练不收敛或准确率上不去:梯度、学习率和数据问题
MNIST本身是个很稳的入门任务,出现严重不收敛的情况通常有三个原因。
学习率不合适是首因。Adam的默认学习率1e-3在大多数情况下可用,但如果你自己改成了1e-1,损失很可能直接变成nan。学习率太大会让参数更新跨过最优点,在损失曲面上震荡;学习率太小会让收敛速度慢到看起来像没训练。判断方法很简单:看前几个iteration的loss变化。如果loss在几百个step内从2.3附近持续下降,说明学习率合理;如果loss上下剧烈波动,先调小学习率10倍再说。
初始化问题也常见。如果模型所有参数初始值一样,或者偏置初始值不合适,梯度更新会同步进行,网络学不到有区分度的特征。PyTorch默认的初始化方式对CNN已经做了合理配置,所以这个项目里不需要手动干预。但如果自己用nn.Parameter手工初始化权重,最好用较小的随机数,避免激活值在ReLU前全部为负导致梯度消失。
数据问题最容易被忽略。MNIST的数据标签是确定的,但如果你拆分数据集时候写错了,比如训练集里混入测试集,或者transform参数顺序不对,模型学到的规律就是错的。这个项目里每次要把完整流程跑通一遍,建议先从几十张图上做小规模冒烟测试,batch_size=4、跑一个epoch,确认损失有下降、维度没问题,再上完整数据集。
5.3 过拟合与提升泛化能力:从99.0%往99.5%迈进的微调方向
模型跑到99%测试准确率之后,想继续往上提升,靠"加更多层"收益不大。这个阶段要考虑的是泛化能力的精细化调整。
我实测有效的一个操作是数据增强。MNIST虽然图像简单,但加入RandomAffine做轻微旋转和平移,可以让模型对笔迹偏移更鲁棒。具体来说,transforms.RandomAffine(degrees=10, translate=(0.1, 0.1))在训练集上做,测试集不做任何增强,通常能把测试准确率从99.0%提升到99.3%左右。这里要强调,测试集永远不能加随机增强,因为评估的是模型在干净数据上的表现。
另一个操作是调整优化器。训练后期把Adam换成SGD加动量,或者对Adam设置weight_decay=1e-4做L2正则化,能限制权重过大,提升泛化。L2正则化的思想很直观:让权重整体保持较小的值,避免网络在训练集上过度拟合噪声。加入正则化之后训练准确率可能微降,但测试准确率通常能涨一点。
还可以试试调整通道数。把conv1从32换成64,conv2从64换成128,参数量增加四倍多,但测试准确率提升有限。这说明MNIST任务简单,模型容量早已超出需求。这也提醒我们一个重要的工程观念:决定模型上限的不只是参数量,结构设计和数据质量同等重要。不是模型越大越好,而是在满足任务需求的前提下,把每一分计算都花在刀刃上。
5.4 训练效率优化:从CPU到GPU的平滑迁移
如果你现在只有CPU,跑完10个epoch大概需要5到10分钟,这个速度完全可以接受。但如果想用GPU加速,代码里已经写好了设备判断逻辑,torch.device('cuda' if torch.cuda.is_available() else 'cpu')之后,把模型和每个batch的数据都.to(device)。整个过程不用改任何其他代码。
这里有一个容易踩的坑:模型在CPU上创建后,如果先跑了几个epoch,再用.to(device)迁移到GPU,Adam优化器里的状态可能需要重置,否则可能报错或者表现异常。稳妥的做法是一开始在代码里就做设备判断,让模型自始至终在同一个设备上训练。另外,数据加载时num_workers参数可以加快CPU到GPU的数据传输,但对MNIST这种小数据集,默认设置就够了,不用刻意调优。
训练时间上,一个现代化的GPU跑MNIST一个epoch通常只需要几秒钟,10个epoch一分钟之内就能完成。这时候瓶颈反而不在计算,而在数据加载。如果数据已经在内存里,DataLoader可以直接用from torch.utils.data import DataLoader配合batch_size=64,加载开销很小。真正到了大规模数据集时,才需要考虑num_workers和预取策略。
我个人在训练效率这件事上的体会是:MNIST阶段不要花太多时间调性能,它的规模决定了一切优化收益都很有限。真正有用的,是养成几个好习惯——设备判断写早一点、梯度清空不遗漏、评估时关梯度、数据放到device上再进模型。这些小习惯在大项目里会比任何技巧都更早救你命。把MNIST当作这些习惯的练习场,而不是性能调优的战场,这是我觉得做这个项目最划算的收获。