文章目录一、MNIST数据集简介二、完整代码实现三、核心模块讲解3.1 数据集与DataLoader3.2 设备自动选择3.3 全连接神经网络模型搭建3.4 训练函数train()完整训练五步3.5 测试函数test()模型评估3.6 损失函数与优化器3.7 PyTorch常用损失函数一览四、运行结果说明五、常见踩坑总结六、总结一、MNIST数据集简介MNIST手写数字数据集是深度学习入门最经典数据集一共70000张28×28单通道灰度图片训练集60000张用来训练神经网络权重测试集10000张用来评估模型泛化能力每张图片对应标签0‑9代表手写数字类别像素原始范围0‑255代码中转换为张量后归一化到0~1之间。二、完整代码实现 MNIST手写数字数据集介绍 一共70000张灰度图片60000张训练集10000张测试集。 图片大小28×28像素单通道灰度图。 fromtorchimportnn# 导入神经网络模块fromtorch.utils.dataimportDataLoader# 数据加载器把数据集分批次打包fromtorchvisionimportdatasets# torchvision内置数据集库MNIST在这里fromtorchvision.transformsimportToTensor# 转换器PIL图片 → PyTorch张量Tensor构建训练数据集对象training_datadatasets.MNIST(rootdata,# 数据集本地存放文件夹在当前项目目录下生成data文件夹trainTrue,# True代表读取训练集(6万张)downloadTrue,# True本地没有文件就联网下载已有文件直接跳过下载transformToTensor(),# 将图片转为Tensor张量像素0~255归一化到0.0~1.0)构建测试数据集对象test_datadatasets.MNIST(rootdata,# 和训练集存到同一个data目录trainFalse,# False读取测试集(1万张)用来评估模型效果downloadTrue,transformToTensor(),)importmatplotlib.pyplotasplt# 创建画布可视化查看9张手写数字图片figureplt.figure()foriinrange(9):img,labeltraining_data[i]# 取出第i张图片和对应的数字标签(0~9)figure.add_subplot(3,3,i1)# 创建3行3列子图依次摆放图片plt.title(label)# 子图标题显示真实数字标签plt.axis(off)# 关闭坐标轴只看图片plt.imshow(img.squeeze(),cmapgray)# squeeze去掉通道维度gray灰度图显示aimg.squeeze()plt.show()# 弹出图片窗口# DataLoader数据集分批次每一批64张图片train_dataloaderDataLoader(training_data,batch_size64)test_dataloaderDataLoader(test_data,batch_size64)# 打印一批数据的shape看懂数据维度forX,yintest_dataloader:print(fShape of X [N, C, H, W]:{X.shape})# N批次大小C通道H高W宽print(fShape of y:{y.shape}{y.dtype})# y标签的shape和数据类型break# 判断设备优先cuda(GPU)苹果设备mps都没有就用cpudevicecudaiftorch.cuda.is_available()elsempsiftorch.backends.mps.is_available()elsecpuprint(fUsing{device}device)# 自定义神经网络类继承nn.ModuleclassNeuralNetwork(nn.Module):def__init__(self):super().__init__()# 调用父类nn.Module的初始化self.flattennn.Flatten()# Flatten展平层把28*28图片拉直成一维向量784self.hidden1nn.Linear(28*28,128)# 第一层全连接输入784输出128神经元self.hidden2nn.Linear(128,256)# 第二层全连接输入128输出256神经元self.outnn.Linear(256,10)# 输出层输入256输出10个类别(数字0‑9)# 前向传播定义数据流动路线函数名forward固定defforward(self,x):xself.flatten(x)# 将图片展平xself.hidden1(x)# 第一层全连接计算xtorch.sigmoid(x)# sigmoid激活函数引入非线性xself.hidden2(x)# 第二层全连接计算xtorch.sigmoid(x)# sigmoid激活函数xself.out(x)# 输出层得到10个类别的预测分数returnx# 创建模型对象并迁移到GPU/CPU设备上modelNeuralNetwork().to(device)print(model)# 打印网络结构 训练函数 dataloader训练数据加载器 model神经网络模型 loss_fn损失函数 optimizer优化器 deftrain(dataloader,model,loss_fn,optimizer):model.train()# 设置模型为训练模式开启dropout等训练专属逻辑本网络没用dropout但规范写法保留batch_size_num1# 记录当前是第几个batchforX,yindataloader:# 将图片数据、标签都搬运到GPU/CPU设备X,yX.to(device),y.to(device)predmodel.forward(X)# 前向传播得到预测结果lossloss_fn(pred,y)# 计算预测值与真实标签之间的损失optimizer.zero_grad()# 梯度清零上一轮的梯度要清空避免累加loss.backward()# 反向传播自动求各个参数的梯度optimizer.step()# 根据梯度更新神经网络权重w、bloss_valueloss.item()# 把tensor类型loss取出普通python数值ifbatch_size_num%1000:# 每100个batch打印一次损失print(floss:{loss_value:7f}[number:{batch_size_num}])batch_size_num1# 测试函数预留后续写测试、计算准确率逻辑deftest(dataloader,model,loss_fn):sizelen(dataloader.dataset)# 获取测试集总样本数量num_batcheslen(dataloader)# 获取测试集batch打包总个数model.eval()# 设置模型为评估模式停止权重更新test_loss,correct0,0# 初始化测试损失、正确样本计数withtorch.no_grad():# 关闭梯度计算不做反向传播节省显存forX,yindataloader:# 遍历测试集每一个批次X,yX.to(device),y.to(device)# 数据、标签迁移到GPU/CPUpredmodel.forward(X)# 前向传播得到预测输出test_lossloss_fn(pred,y).item()# 累加本批次损失correct(pred.argmax(1)y).type(torch.float).sum().item()# 统计本批次预测正确样本数a(pred.argmax(1)y)# 布尔张量预测是否等于真实标签b(pred.argmax(1)y).type(torch.float)# 将布尔值转为float(1.0/0.0)test_loss/num_batches# 计算测试集平均损失correct/size# 计算测试集整体准确率print(fTest result: \n Accuracy:{(100*correct)}%, Avg loss:{test_loss})loss_fnnn.CrossEntropyLoss()# 交叉熵损失函数多用于多分类任务optimizertorch.optim.SGD(model.parameters(),lr0.01)# SGD随机梯度下降优化器学习率0.01epochs10# 设置训练总轮数完整遍历训练集10次fortinrange(epochs):print(fEpoch{t1}\n-------------------------------)# 打印当前是第几轮训练train(train_dataloader,model,loss_fn,optimizer)# 执行一轮训练更新网络权重print(Done!)# 全部轮次训练完成提示test(test_dataloader,model,loss_fn)# 在测试集上评估模型效果计算loss和准确率三、核心模块讲解3.1 数据集与DataLoaderdatasets.MNISTtorchvision内置数据集downloadTrue自动下载数据集到本地data文件夹。ToTensor()把图片转为张量像素值归一化0‑1。DataLoader对数据集做分批次(batch)支持打乱、多线程读取。本例batch_size64每次给模型喂64张图片。数据维度格式[N, C, H, W]Nbatch批次大小C通道数H图片高度W图片宽度。MNIST灰度图C1。3.2 设备自动选择devicecudaiftorch.cuda.is_available()elsempsiftorch.backends.mps.is_available()elsecpu自动优先使用NVIDIA GPU(cuda)苹果硅芯片MPS最后降级CPU。模型和张量必须.to(device)搬运到对应设备才能运算。3.3 全连接神经网络模型搭建继承nn.Module是PyTorch自定义网络标准写法。nn.Flatten()将28×28图片展平成784维一维向量。nn.Linear全连接层实现y x W b yxWbyxWb。forward()函数必须定义描述数据前向流动路径不要手动调用模型对象(X)会自动调用forward。torch.sigmoid()激活函数引入非线性没有激活函数多层网络等价于单层线性模型。网络结构Flatten(784) → Linear(784→128) → sigmoid → Linear(128→256) → sigmoid → Linear(256→10)输出10维向量代表数字0‑9各个类别的得分。3.4 训练函数train()完整训练五步深度学习训练循环五大步骤前向传播pred model(X)得到预测输出计算损失loss loss_fn(pred,y)对比预测与真实标签差距梯度清零optimizer.zero_grad()梯度会累加每轮必须清空反向传播求梯度loss.backward()自动计算所有权重梯度参数更新optimizer.step()使用梯度更新w、b权重model.train()训练模式部分层(Dropout、BN)会启用训练逻辑。3.5 测试函数test()模型评估model.eval()评估模式关闭dropout、batchnorm训练行为。with torch.no_grad()关闭梯度计算节省内存测试阶段不需要反向传播。pred.argmax(1)取10维输出分数最大的下标即为预测数字类别。统计correct预测正确样本数量除以总样本得到准确率。3.6 损失函数与优化器损失函数 CrossEntropyLoss多分类任务首选内部集成LogSoftmaxNLLLoss输出层不需要额外加softmax。优化器 SGD随机梯度下降lr0.01为学习率控制每一步权重更新幅度。Epoch一轮epoch代表完整遍历全部训练集一次本例设置10轮完整训练。3.7 PyTorch常用损失函数一览损失函数使用场景CrossEntropyLoss多分类BCEWithLogitsLoss二分类MSELoss回归任务NLLLoss配合LogSoftmax多分类SmoothL1Loss/HuberLoss回归抗异常值四、运行结果说明程序运行首先自动下载MNIST数据集到./data文件夹。弹出matplotlib窗口展示9张手写数字样本。控制台打印张量维度、使用设备、网络结构。训练过程每100个batch打印loss损失值正常loss会逐步下降。10轮epoch训练结束后执行test函数输出测试集准确率和平均损失。提示如果使用sigmoid激活的简单全连接网络MNIST准确率一般可以达到95%左右。想要更高准确率可以改用ReLU激活、CNN卷积网络。五、常见踩坑总结忘记把model、X、y搬运到deviceCPU/GPU张量混合报错。训练循环忘记optimizer.zero_grad()梯度累加loss不收敛。测试阶段忘记model.eval()和torch.no_grad()显存占用高、评估结果异常。CrossEntropyLoss使用时自己额外加Softmax层会导致效果变差。forward不要手动调用model.forward(X)规范写法是model(X)。六、总结到这里我们就完整跑通了使用 PyTorch 解决 MNIST 手写数字识别的全流程。虽然这是一个入门项目但它涵盖了深度学习开发最核心的几个环节数据流转从 datasets 加载到 DataLoader 分批我们掌握了处理图像数据的标准姿势特别是 [N, C, H, W] 这个维度的概念以后处理任何视觉任务都离不开它。模型构建通过继承 nn.Module我们搭起了一个包含展平层、全连接层和激活函数的基础网络。这一步让你理解了数据是如何在网络中一层层流动并发生变换的。训练闭环这是最重要的一环。前向传播算预测、计算 Loss、梯度清零、反向传播求导、优化器更新参数——这五步法是深度学习的肌肉记忆必须烂熟于心。避坑与规范我们在代码中实践了设备自动切换GPU/CPU、训练/评估模式切换train/eval以及关闭梯度计算no_grad这些都是写出健壮代码的关键细节。虽然 MNIST 数据集规模较小、全连接网络结构相对简单但它犹如深度学习领域的“Hello World”麻雀虽小五脏俱全。掌握了这套标准训练范式你就具备了迁移学习的能力了。