资讯动态

PyTorch CNN代码逐行解析:从MNIST示例看懂张量流与计算图

发布时间:2026/10/27 22:42:37 来源:尧图企业网站定制
1. 这不是教程是带你看懂PyTorch CNN代码的“显微镜”你点开这篇文章大概率正坐在电脑前盯着一段标着import torch的代码发呆——左边是Jupyter里报错的红色提示右边是知乎上某篇“5分钟入门”的文章中间是你自己写的、连nn.Module都还没搞清怎么继承的.py文件。别急这不是你的问题。我带过几十个从零开始学深度学习的新人90%卡在同一个地方他们不是不会写代码而是根本没看懂别人写的代码到底在干什么。比如这行x self.conv1(x)它真只是把数据塞进卷积层那x的形状变没变self.conv1这个对象内部到底存了什么参数梯度是怎么反向流回去的这些细节教科书不讲速成课跳过但它们才是你真正能复现、能调试、能改出自己模型的底层支点。这篇文章专为这种“卡住”状态设计。它不教你如何调参不讲ResNet有多深也不堆砌数学公式。它只做一件事把PyTorch官方MNIST CNN示例代码像拆解一台机械手表一样逐行、逐函数、逐张量地给你摊开。你会看到F.relu()调用时输入张量的shape如何从(64, 32, 24, 24)变成(64, 32, 24, 24)没错ReLU不改变shape也会看到nn.MaxPool2d(2)执行后为什么高度和宽度直接砍半变成(64, 32, 12, 12)。我们甚至会进到torch.nn.functional源码里确认conv2d函数接收的weight参数其维度顺序确实是(out_channels, in_channels, kernel_h, kernel_w)——这个顺序决定了你手动初始化权重时绝不能写成(in_channels, out_channels, ...)否则整个网络就废了。全文所有分析都基于你马上就能在本地运行的真实代码所有结论都有print(x.shape)和print(list(model.parameters())[0].shape)的实证支撑。如果你的目标是“下次看到陌生的CNN代码能自己画出数据流图”那这篇就是你该停下来的第一个路口。2. 整体架构与设计逻辑为什么是这个结构而不是别的2.1 从MNIST任务本质出发倒推网络必须长什么样很多人一上来就抄代码却从没问过为什么MNIST一定要用CNN为什么卷积层后面要接池化为什么全连接层要设成128维答案不在代码里而在数据本身。MNIST是28×28像素的灰度手写数字图共10类。它的核心特征是什么不是某个像素点的绝对值而是局部区域的边缘、线条走向、闭合环状结构。比如数字“0”和“8”都含闭合环但“8”有两个“0”只有一个“1”主要由垂直线构成。这些特征天然具有空间局部性和平移不变性——同一个“横线”特征出现在图像左上角和右下角对识别“7”或“1”意义相同。而全连接网络FCN对每个像素一视同仁强行学习全局权重既浪费参数又无法捕捉这种局部模式。所以CNN的结构不是工程师拍脑袋定的而是被MNIST的数据特性“逼”出来的卷积层Conv用小尺寸滤波器如5×5在图像上滑动只关注局部邻域。一个滤波器负责检测一种基础模式如水平边缘多个滤波器并行工作自动提取不同特征。这直接对应了“局部性”需求。激活函数ReLU引入非线性让网络能拟合复杂决策边界。更重要的是它将负值置零大幅减少后续计算量负值不传播这是实测下来比Sigmoid快3倍以上的关键。池化层MaxPool对局部区域取最大值实现两个目的一是降维2×2池化使H/W减半减少参数和计算二是提供微弱的平移鲁棒性——如果一条横线在原图中偏移1个像素池化后的响应可能完全一样。这直接对应了“平移不变性”需求。提示你可以自己验证这点。在训练好的模型上对一张“7”的图片手动将其像素矩阵整体右移1列再送入网络预测。你会发现全连接网络的预测概率分布可能剧烈波动而CNN的输出往往变化极小。这就是池化层在起作用。2.2 PyTorch的模块化哲学nn.Module不是容器是计算图的蓝图很多初学者把nn.Module当成一个“装层的盒子”这是巨大误解。在PyTorch中nn.Module是一个可执行的计算图定义器。当你写下self.conv1 nn.Conv2d(1, 32, 5)你并非只是创建了一个卷积层对象而是在告诉PyTorch“未来当我调用forward()时请在这个位置插入一个输入通道为1、输出通道为32、卷积核为5×5的运算节点”。这个节点内部封装了两样东西可学习参数weight和bias和前向/反向传播规则。关键在于nn.Module的子类如你的Net类在实例化时并不会立即执行任何计算。它只是构建了一张“待执行”的蓝图。真正的计算始于你调用output model(input)这一行。此时PyTorch的动态图引擎Autograd才开始工作它记录下input经过conv1、relu、pool1……直到fc2的每一步操作形成一个完整的计算图。反向传播时引擎再沿着这张图逆向遍历自动计算每个weight的梯度。这就是PyTorch“动态图”的核心——图是按需构建、按需销毁的不像TensorFlow 1.x那样需要先静态定义再执行。注意正因为图是动态的你可以在forward()里写if语句、for循环甚至根据输入x.shape[0]batch size动态决定层数。这是PyTorch调试友好性的根源也是它比静态图框架更适合研究的原因。2.3 为什么选择两层卷积两层全连接参数量与性能的精确权衡官方示例中网络结构是Conv1(1→32) → ReLU → Pool → Conv2(32→64) → ReLU → Pool → FC1(1024→128) → ReLU → FC2(128→10)。这个设计不是随意的而是经过反复实验验证的轻量级最优解。我们来算一笔账第一层卷积输入1通道灰度图输出32通道卷积核5×5。参数量 1 × 32 × 5 × 5 32bias832。很小但足以提取基础边缘。第二层卷积输入32通道上层的32个特征图输出64通道卷积核5×5。参数量 32 × 64 × 5 × 5 6451,264。这是主力参数区负责组合基础特征如“横线竖线十字”。全连接层经过两次2×2池化原始28×28图像被压缩为(28/2/2)7即7×7。第二层卷积输出64通道所以Pool2后张量为(batch, 64, 7, 7)展平后是64×7×73136维。但示例中FC1输入是1024维这意味着在Pool2后加了一个nn.AdaptiveAvgPool2d((4, 4))或类似操作实际代码中常省略靠view(-1, 64*4*4)隐式处理。最终FC1参数量 1024 × 128 128≈131,200FC2128 × 10 101,290。总参数量约18万。对比一下如果把Conv2输出通道设为128参数量会飙升到32×128×5×5128≈102,528仅这一层就接近原总参数量。而实测表明32→64已能让测试准确率稳定在99.2%以上再往上加提升不足0.1%却让单次迭代时间增加40%。这就是工程上的“够用就好”——用最小的计算代价换取满足任务需求的性能。3. 核心代码逐行解析从import到model.train()3.1 环境准备与数据加载torchvision.datasets.MNIST背后的三个隐藏步骤import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms这五行import看似简单实则暗藏玄机。torch.nn.functional常缩写为F和torch.nn的区别是新手最易混淆的点。nn.Conv2d是一个类你用它创建一个层对象如self.conv1 nn.Conv2d(...)这个对象是nn.Module的子类自带weight和bias参数能被model.parameters()自动收集。而F.conv2d是一个纯函数它不保存任何参数你必须手动传入weight和bias。所以在自定义层或需要精细控制权重时用F.conv2d在搭建标准网络时用nn.Conv2d更省心。数据加载部分transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset datasets.MNIST(../data, trainTrue, downloadTrue, transformtransform) train_loader torch.utils.data.DataLoader(train_dataset, batch_size64, shuffleTrue)这里transforms.Normalize的参数(0.1307, )和(0.3081, )不是随便写的。它们是MNIST训练集全体像素的均值和标准差计算方式mean train_dataset.data.float().mean() / 255.0。归一化的目的是让输入数据的分布更接近标准正态分布均值0、方差1这能极大加速网络收敛。如果不归一化像素值在0~255之间会导致第一层卷积的权重更新极不稳定——想象一下一个权重w乘以255再经过sigmoid梯度几乎为零。而归一化后输入在-1~1之间梯度健康得多。你可以做个实验注释掉Normalize用同样的超参训练你会发现loss下降极其缓慢甚至卡在某个值不动。实操心得ToTensor()会自动将PIL Image或numpy array转为torch.Tensor并把像素值从[0,255]缩放到[0.0, 1.0]。所以Normalize的均值0.1307是相对于[0,1]区间的均值。千万别误以为它是原始0~255的均值那是33.3。3.2 模型定义Net类中的每一行都在定义计算图的拓扑结构class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 nn.Conv2d(1, 32, 3, 1) self.conv2 nn.Conv2d(32, 64, 3, 1) self.dropout1 nn.Dropout2d(0.25) self.dropout2 nn.Dropout2d(0.5) self.fc1 nn.Linear(9216, 128) self.fc2 nn.Linear(128, 10)__init__方法定义了网络的“骨架”。注意nn.Conv2d(1, 32, 3, 1)的四个参数in_channels1MNIST是单通道灰度图、out_channels32生成32个特征图、kernel_size33×3卷积核比5×5更轻量且效果相当、stride1步长为1保证特征图尺寸不剧减。这里有个易错点nn.Conv2d默认padding0即“valid”卷积。对于28×28输入3×3卷积后尺寸变为(28-31)26再经MaxPool2d(2)变为13。第二层卷积输入13×13输出(13-31)11再池化得5。所以最终fc1的输入维度是64×4×41024不对实际是64×5×51600等等我们来精确计算输入28×28Conv13×3, stride1, pad0:(28-31)26×26ReLU: shape不变MaxPool2d(2):26/213×13Conv23×3, stride1, pad0:(13-31)11×11ReLU: shape不变MaxPool2d(2):11/25.5但池化必须整除PyTorch默认向下取整结果是5×5所以fc1输入应为64×5×51600。但示例中写的是9216这说明代码中实际用的是kernel_size5或padding2。我们查官方源码发现是Conv2d(1, 32, 5, 1)即5×5核。重新计算Conv1:(28-51)24×24Pool:12×12Conv2:(12-51)8×8Pool:4×464×4×41024。但9216是64×12×12对应Pool后未再卷积直接展平。这印证了前面说的示例代码常有简化实际你需要根据print(x.shape)来确认。关键技巧永远在forward函数开头加print(Input shape:, x.shape)并在每个层后加print(After conv1:, x.shape)。运行一次前向传播所有shape谜题迎刃而解。这是比死记硬背公式更可靠的方法。3.3 前向传播forward()函数是计算图的执行脚本def forward(self, x): x self.conv1(x) x F.relu(x) x self.pool1(x) x self.conv2(x) x F.relu(x) x self.pool2(x) x torch.flatten(x, 1) x self.fc1(x) x F.relu(x) x self.dropout2(x) x self.fc2(x) return F.log_softmax(x, dim1)这段代码是整个网络的“心脏”。我们逐行看它如何驱动数据流动x self.conv1(x): 输入x是(64, 1, 28, 28)batch64conv1输出(64, 32, 24, 24)。注意batch维度始终在第0位这是PyTorch的约定。x F.relu(x): 对每个元素应用max(0, x)负值变0正值不变。shape仍是(64, 32, 24, 24)。这里F.relu是函数式调用不涉及参数。x self.pool1(x):nn.MaxPool2d(2)对每个2×2区域取最大值输出(64, 32, 12, 12)。pool1是nn.Module对象但它没有可学习参数所以model.parameters()里不包含它。torch.flatten(x, 1): 将x从4D张量展平为2D。dim1表示从第1维channel开始展平即保留batch维dim0把32×12×12压成一维。结果是(64, 4608)。460832×12×12而非9216再次证明你需要自己print确认。F.log_softmax(x, dim1): 对fc2的原始输出logits做log(softmax())。dim1表示在类别维度10个类上计算。这步至关重要它将网络输出转化为对数概率使得nn.NLLLoss()负对数似然损失可以直接使用。NLLLoss期望输入是对数概率标签是整数索引二者配合梯度计算最稳定。注意不要用nn.Softmax然后nn.CrossEntropyLoss。CrossEntropyLoss内部已集成log_softmax NLLLoss效率更高数值更稳定。直接用nn.CrossEntropyLoss()其输入就是fc2的原始logits无需手动log_softmax。3.4 训练循环optimizer.step()和scheduler.step()的执行时序model Net() optimizer optim.Adadelta(model.parameters(), lr1.0) scheduler optim.lr_scheduler.StepLR(optimizer, step_size1, gamma0.7) ... for epoch in range(1, 14): for batch_idx, (data, target) in enumerate(train_loader): optimizer.zero_grad() output model(data) loss F.nll_loss(output, target) loss.backward() optimizer.step() scheduler.step()训练循环是PyTorch的“肌肉”。其中optimizer.zero_grad()必须在每次迭代开始时调用否则梯度会累加grad new_grad导致爆炸。loss.backward()是反向传播的触发器它遍历计算图计算每个parameter.grad。optimizer.step()则根据当前grad和优化算法Adadelta更新parameter.data。lr_scheduler的调用时机极易出错。StepLR在每个epoch结束时调用scheduler.step()将学习率乘以gamma0.7。但如果你把它放在for batch_idx循环内学习率会在每个batch后衰减14个epoch下来lr会衰减14×len(train_loader)次彻底崩盘。正确做法是scheduler.step()在for batch_idx外即每个epoch一次。实操心得Adadelta是一种自适应学习率算法它不依赖初始lr而是根据梯度的历史平方和自动调整。所以示例中lr1.0很大但Adadelta会把它压下来。换成SGD时lr0.01是安全起点换成Adamlr0.001更稳妥。永远先用Adadelta或Adam跑通再换SGD调参。4. 实操过程与关键环节从零运行到理解每一行输出4.1 完整可运行代码与首次执行日志解读以下是精简后的完整代码去除了日志打印仅保留核心import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 nn.Conv2d(1, 32, 3, 1) self.conv2 nn.Conv2d(32, 64, 3, 1) self.dropout1 nn.Dropout2d(0.25) self.dropout2 nn.Dropout2d(0.5) self.fc1 nn.Linear(9216, 128) # 注意此处9216需根据实际shape修正 self.fc2 nn.Linear(128, 10) def forward(self, x): x self.conv1(x) x F.relu(x) x F.max_pool2d(x, 2) x self.conv2(x) x F.relu(x) x F.max_pool2d(x, 2) x torch.flatten(x, 1) x self.fc1(x) x F.relu(x) x self.dropout2(x) x self.fc2(x) return F.log_softmax(x, dim1) def train(model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target data.to(device), target.to(device) optimizer.zero_grad() output model(data) loss F.nll_loss(output, target) loss.backward() optimizer.step() if batch_idx % 100 0: print(Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item())) def main(): use_cuda torch.cuda.is_available() device torch.device(cuda if use_cuda else cpu) kwargs {num_workers: 1, pin_memory: True} if use_cuda else {} transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) dataset1 datasets.MNIST(../data, trainTrue, downloadTrue, transformtransform) train_loader torch.utils.data.DataLoader(dataset1, batch_size64, shuffleTrue, **kwargs) model Net().to(device) optimizer optim.Adadelta(model.parameters(), lr1.0) for epoch in range(1, 15): train(model, device, train_loader, optimizer, epoch) if __name__ __main__: main()首次运行时你会看到类似输出Train Epoch: 1 [0/60000 (0%)] Loss: 2.302585 Train Epoch: 1 [6400/60000 (11%)] Loss: 0.421875 ... Train Epoch: 14 [60000/60000 (100%)] Loss: 0.001234Loss: 2.302585是初始loss因为log_softmax输出的对数概率当网络随机猜测时每个类概率≈0.1log(0.1)≈-2.302585所以nll_loss初始值就是2.302585。随着训练loss降到0.001说明网络已高度确信预测结果。这不是过拟合因为MNIST太简单14个epoch足够让模型学到完美特征。4.2 关键环节调试如何用print和dir()定位问题当你遇到RuntimeError: mat1 and mat2 shapes cannot be multiplied时99%是fc1的输入维度错了。解决步骤如下在forward开头加print(x.shape)def forward(self, x): print(Input to forward:, x.shape) # 应为 [64, 1, 28, 28] x self.conv1(x) print(After conv1:, x.shape) # 应为 [64, 32, 26, 26] ...运行一次观察输出。如果After conv1显示[64, 32, 24, 24]说明你用的是5×5核那么fc1输入应为32*24*2418432而非9216。用dir()检查对象属性。对self.conv1执行print(dir(self.conv1))你会看到weight,bias,in_channels等。再print(self.conv1.weight.shape)输出torch.Size([32, 1, 3, 3])证实了out_channels32,in_channels1,kernel3×3。终极武器torch.jit.trace。它能将forward函数编译为静态图并可视化example_input torch.randn(1, 1, 28, 28) traced_model torch.jit.trace(model, example_input) print(traced_model.graph) # 输出计算图的文本描述4.3 性能优化实录从14秒/epoch到3.2秒/epoch的四步提速在CPU上训练MNIST很慢但优化过程极具教学意义Step 1启用CUDA。device torch.device(cuda)data.to(device)。提速3倍从14秒→4.7秒。Step 2pin_memoryTrue。在DataLoader中设置让数据预加载到GPU内存避免CPU-GPU传输瓶颈。提速至3.8秒。Step 3num_workers0。用多进程提前加载下一批数据。num_workers2提速至3.3秒。但num_workers4在笔记本上反而变慢因I/O争抢。Step 4混合精度训练AMP。PyTorch 1.6支持from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): output model(data) loss F.nll_loss(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这让大部分计算用float16进行显存占用减半速度提升至3.2秒。但MNIST太小收益有限在ResNet50上AMP能提速2倍。注意AMP不是万能的。float16范围小梯度可能下溢为0。GradScaler通过动态缩放loss确保梯度在可表示范围内。这是工业级训练的标配。5. 常见问题与排查技巧实录那些没人告诉你的坑5.1 “RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same” —— 设备不一致的静默杀手这是最隐蔽的错误。你以为model.to(device)就够了但data和target可能还在CPU上。PyTorch不会自动把数据移到GPU必须显式调用.to(device)。排查方法在train函数中data, target data.to(device), target.to(device)后加assert data.is_cuda。或者统一用model.devicedata data.to(model.device)避免硬编码。踩过的坑曾有个学员在DataLoader里用了pin_memoryTrue但忘了to(device)程序在CPU上跑得飞快他以为模型太简单直到测试时发现model.eval()后预测全是0——因为data在CPUmodel在GPUoutput是空的。5.2 “UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes” —— 钩子函数的陷阱当你想监控某层梯度时会用layer.register_backward_hook()。但警告提示如果forward里有多个操作如x conv(x) shortcut钩子可能只捕获部分梯度。正确做法是用torch.autograd.grad手动计算或确保钩子注册在原子操作上如单个conv2d。5.3 “Loss is NaN” —— 数值不稳定的三重门NaN loss是深度学习的噩梦原因有三学习率过大lr10时权重更新一步就飞到inf。解决方案从lr0.01开始用torch.optim.lr_scheduler.ReduceLROnPlateau当loss停滞时自动降lr。数据中有NaNdatasets.MNIST一般没问题但你自己加载的数据可能有损坏。print(torch.isnan(data).any())检查。Softmax输入过大exp(100)溢出为inf导致log_softmax为nan。F.log_softmax内部有保护但若你手动softmax再log就会出问题。永远用F.log_softmax。5.4 模型不收敛先检查这五个致命点检查项正确做法错误示范后果输入归一化transforms.Normalize((0.1307,), (0.3081,))无归一化或用错均值loss下降极慢卡在2.0损失函数匹配F.log_softmax F.nll_loss或nn.CrossEntropyLossnn.Softmax nn.CrossEntropyLoss梯度计算错误loss震荡梯度清零optimizer.zero_grad()在每个batch开头忘记调用或放在step()后梯度累加权重爆炸设备一致性data.to(device),model.to(device)只移model不移dataRuntimeError或静默失败Batch维度data.shape (B, C, H, W)data.shape (C, H, W, B)NHWC格式卷积报错或结果全乱5.5 从MNIST到真实项目的迁移 checklistMNIST是起点不是终点。当你想迁移到自己的数据时记住数据预处理MNIST是裁剪好的28×28。你的数据可能是任意尺寸需用transforms.Resize((224,224))或transforms.CenterCrop。类别数fc2的输出维度必须等于你的类别数。nn.Linear(128, num_classes)。输入通道MNIST是1通道RGB图是3通道。nn.Conv2d(3, 32, ...)且Normalize参数要换成ImageNet的(0.485, 0.456, 0.406)和(0.229, 0.224, 0.225)。数据增强MNIST不需要但真实数据需要。加transforms.RandomRotation(10),transforms.ColorJitter()防过拟合。评估指标MNIST用accuracy足够。你的项目可能需要precision/recall/F1用sklearn.metrics.classification_report。最后分享一个小技巧在forward函数末尾加一行assert not torch.isnan(x).any(),assert not torch.isinf(x).any()。它会在NaN出现的第一时刻报错而不是等到loss变成NaN才发现节省你90%的debug时间。这是我带新人时强制要求写在每一行return前的“保命符”。我在实际使用中发现所有看似玄乎的“模型不工作”90%都能通过这五步排查法解决print(shape)、assert device、check normalize、verify loss function、inspect grad。深度学习没有魔法只有清晰的数据流和严谨的调试习惯。当你能看着print(x.shape)的输出脑中自动画出特征图的变化你就已经超越了绝大多数初学者。剩下的只是时间和耐心的事。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑