PyTorch的torch.nn模块是我用得最多的一个工具箱，没有之一。说实话，学深度学习如果只死啃理论，不亲手把神经网络搭出来跑一遍，很多概念是飘着的。而一旦进了torch.nn这个工具箱，从网络层搭建、损失函数选型到优化器配置、模型保存加载，你几乎每天都在跟它打交道。本文就围绕第3章「PyTorch神经网络工具箱」这套内容，把我在实际开发中反复用到、也反复踩坑的要点梳理一遍。适合刚学完张量操作、准备进入神经网络实战阶段的读者参考，也适合已经跑过Demo但想搞清楚内部工作机制的开发者对照自查。

1. nn.Module：神经网络的基本组织单位，读懂它就读懂了一半PyTorch

1.1 为什么所有网络层都必须继承nn.Module

初学者最容易出现的一个问题是：我能不能自己写个类，不用继承nn.Module，只要在forward里做矩阵乘法就行？技术上可以，但你会立刻发现丢掉了很多基础设施。nn.Module这个基类干的事情很关键——它把参数管理、子模块管理、设备切换、序列化和训练/评估切换这些底层能力全部打包好了。你写一个自定义层，只需要继承它，然后把可训练参数声明成nn.Parameter，PyTorch就会自动把参数注册进模型的parameter列表里。

我自己第一次手写一个两层感知机时，就偷懒没用nn.Module继承，直接在普通类里用了torch.Tensor当成权重。训练的时候发现无论怎么调用optimizer.step()，模型的权重死活不更新。查了半天才意识到，optimizer需要从model.parameters()拿到参数引用，而普通的Tensor属性根本不会被注册，optimizer手里是空的，当然不会更新。这个坑几乎每个新手都会踩一次。后来我形成习惯：凡是可训练参数，一律用nn.Parameter，或者直接用nn.Linear这类封装好的层，让它们内部去注册参数。

1.2 forward方法：自定义前向传播的逻辑入口

nn.Module的调用方式是model(x)而不是model.forward(x)。这一点经常有人搞混，我自己早期也犯过同样的错误。在PyTorch里，调用model(input)时实际执行的是__call__方法，而__call__内部会依次处理hooks、检查参数状态，然后再调用forward。如果你直接调用model.forward(input)，等于绕过了整个Module的预处理和后处理机制，比如模型里挂了forward hook的话，hook就不会被触发。

在自定义forward时，需要注意一个核心原则：尽量让每一层都能接受[batch, ...]格式的输入。因为训练时数据通常是一个batch一起传入，前向传播在batch维度上并行计算。如果你的层在forward里不小心用了循环逐样本处理，虽然能跑，但速度会慢得让人抓狂——因为该并行的矩阵运算被退化成串行。能向量化就向量化，是我写forward函数时的一条铁律。

1.3 parameters、buffers、state_dict：三个容易混淆的存取接口

这三个接口表面上都能返回模型内部状态，但语义完全不同。model.parameters()返回的是所有可训练参数（比如权重、偏置），它们参与梯度更新。model.buffers()返回的是不参与梯度更新的状态量，典型代表是BatchNorm里运行的均值和方差。model.state_dict()则是把参数和buffer一起打包成有序字典，常用于模型保存和加载。

一个实际场景：用BatchNorm时，如果你自己实现保存模型逻辑，只保存parameters而不保存buffers，恢复模型后BatchNorm的running_mean和running_var就会丢失，而模型在推理时的输出会因为统计量不一致而出现明显偏差。所以我的习惯是保存模型一律用torch.save(model.state_dict())，这样参数和buffer一起保存，一网打尽。

注意：加载模型后如果发现预测结果和训练时对不上，优先检查是不是只恢复了权重、丢了BatchNorm的buffer。

1.4 子模块注册：Python list到底是坑还是工具

在自定义一个复杂网络时，你可能很自然地想在__init__里写self.layers = [nn.Linear(10, 10), nn.ReLU()]这行代码，把一堆层放进Python list。问题来了：这样写不会报错，但model.parameters()里根本不会有这些层的参数。因为nn.Module只会自动识别直接赋值成nn.Module或包含nn.Module的container类型属性（比如nn.Sequential、nn.ModuleList），普通list里的子模块不会被跟踪。

我最初构建一个多层堆叠结构时，就遇到过这种情况：网络能前向传播，但loss一直不动。调试了很久才发现，所有线性层的权重压根没进参数列表。换成nn.ModuleList之后就正常了。这个点单独拿出来说，是因为它的报错极其隐蔽——代码不报异常，训练流程看起来也正常，但模型完全学不到东西。

2. 常用网络层的选型思路：Linear、Conv2d、Dropout各管一摊

2.1 nn.Linear：全连接层背后的初始化逻辑

nn.Linear是PyTorch里使用频率最高的层之一。它做的事情就是一个线性变换y = xW^T + b。但很多人没注意到的是，它的参数初始化方式默认是Kaiming均匀分布，范围由输入的fan_in决定。这个初始化策略对接ReLU激活函数很友好，因为它考虑了ReLU会屏蔽一半神经元的情况，避免网络在初期就陷入梯度消失。

如果你在自定义层里用nn.init.kaiming_normal_或xavier_normal_，需要确认自己选择的初始化方式和激活函数是否匹配。我从实践中总结的经验是：ReLU系激活搭配Kaiming初始化，Sigmoid/Tanh搭配Xavier初始化效果更稳。如果你用nn.Linear默认初始化，通常不需要额外操心；但如果你想用自定义初始化方法，记得在__init__里对每个层逐一调用初始化函数，或者用apply方法批量处理。

2.2 nn.Conv2d：padding和输出尺寸计算，不背公式也能算明白

卷积层是计算机视觉模型的主干。nn.Conv2d的核心参数是in_channels、out_channels、kernel_size、stride、padding和dilation。输出尺寸的计算公式是：

output_size = floor((input_size + 2padding - dilation(kernel_size-1) - 1) / stride + 1)

很多初学者看到这个公式就头疼。我的建议是直接用PyTorch的工具函数来计算感受野和输出尺寸，或者写个小函数便于调试。实际项目中，我常用kernel_size=3、stride=1、padding=1的组合，这样卷积不会改变特征图的宽高，只在通道数上做变换——这个设计在搭建深层网络时非常省心，不用每层都去算一遍尺寸。

关于padding还有一个细节：padding="same"在PyTorch中只对stride=1的情况有效。如果你设置了stride=2又期望输出尺寸和输入一样，PyTorch会直接报错。这是和TensorFlow行为不一样的地方，跨框架迁移代码时容易踩。

2.3 Dropout与BatchNorm：训练模式和评估模式的行为差异

Dropout在训练时会随机丢弃一部分神经元输出，起到正则化作用；在评估时应该不丢弃任何神经元，而是把输出按保留概率缩放（PyTorch里dropout在eval模式下默认关闭并自动做缩放处理）。BatchNorm在训练时用当前batch的均值和方差归一化数据，同时更新running_mean和running_var；在评估时则直接用已保存的running统计量。

这正是model.train()和model.eval()存在的根本原因。很多人以为model.eval()只是关掉梯度计算，其实它更重要的作用是切换这些层的行为模式。我用PyTorch训练模型时养成的习惯是：训练循环开头写上model.train()，validation和测试阶段开头写上model.eval()，并且用torch.inference_mode()或torch.no_grad()包住前向传播，减少内存和计算开销。

2.4 激活函数：ReLU是默认选项，但别忽略LeakyReLU和GELU

激活函数在神经网络工具箱里数量很多，但日常使用频率最高的还是ReLU。它实现简单、计算快、能有效缓解梯度消失。缺点也明显：神经单元死亡问题——如果某个神经元的输入长期为负，ReLU输出恒为0，梯度流无法通过，这个神经元就再也没机会恢复了。因此我在层数较多、或使用较大学习率的场景下，会更倾向用LeakyReLU，让负区间保留一个小斜率（默认negative_slope=0.01），避免神经元死亡。

另一个值得关注的是GELU，它在Transformer结构中被大量使用。PyTorch中nn.GELU支持精确计算和近似计算两种模式。如果你在复现BERT类模型，激活函数一定要用GELU；如果只是做常规CNN分类，ReLU系列完全够用。选激活函数没有绝对正确，关键看网络结构和训练稳定性。

3. 损失函数与优化器：训练闭环里的两大支柱

3.1 损失函数怎么选：回归、二分类、多分类各有各的讲究

PyTorch把常用的损失函数都收进了torch.nn中。回归任务我用得最多的是nn.MSELoss，它计算预测值和真实值的均方误差，对异常值比较敏感。如果数据集里存在明显的离群点，我会考虑换nn.SmoothL1Loss，它在误差较小时保持平方增长，误差较大时转为线性增长，梯度更温和，不容易被离群点带着走。

分类任务的默认选项是nn.CrossEntropyLoss。有一个高频混淆点：PyTorch的CrossEntropyLoss内部已经包含了Softmax操作，所以你的网络最后一层不需要额外加nn.Softmax作为输出层；否则相当于做了两次softmax，概率分布会被扭曲。我在刚开始写分类模型时，就因为在全连接层后面加了Softmax再加CrossEntropyLoss，训练出来的模型预测概率图总是怪怪的。后来把Softmax去掉，逻辑回归那种情况的loss变化立刻正常了。

二分类其实也可以用nn.BCEWithLogitsLoss，它把Sigmoid也内置了，输入直接是网络logits，标签是0或1。这一点在多标签分类任务中尤其常用，因为每个标签可以独立做二分类，输出维度等于标签数量。

3.2 优化器：SGD的动量机制与Adam的自适应学习率

优化器选型通常会在torch.optim.SGD和torch.optim.Adam之间纠结。SGD配上动量momentum之后，收敛路径更平滑，泛化能力往往更好，但学习率需要手动调节，训练时间更长。Adam自带自适应学习率，对不同参数分别调整步长，收敛速度快，很多任务上开箱即用，但有时候最终精度会略逊于调好的SGD。

我自己的选择逻辑是：任务需要快速验证、快速出结果时，直接用Adam，lr设置成1e-3是常见起步值。当模型结构和数据准备阶段都已经稳定，想要精调模型达到更好泛化性能时，我会切成SGD+momentum=0.9，并将学习率按经验缩小到原来的十分之一再跑。另外，实际项目中我几乎总是给Adam设置weight_decay=1e-5或1e-4，这相当于L2正则化，能抑制参数过大，对训练稳定性有明显帮助。

提示：pytorch里optimizer和model是两个对象。每次step前别忘了调用optimizer.zero_grad()，否则梯度会在多个batch上累积。

3.3 学习率调度：别从头到尾都用同一个学习率

学习率调度能显著影响训练效果。PyTorch内置了许多调度器，最常用的是torch.optim.lr_scheduler.StepLR、MultiStepLR和ReduceLROnPlateau。StepLR按固定步长衰减，比如每30个epoch把学习率乘以0.1。MultiStepLR则让你指定在第几个epoch发生衰减，比如第60、90个epoch衰减一次，这在ImageNet类数据集的训练配置里非常常见。

ReduceLROnPlateau不同，它不看epoch数，而是监控某个指标（比如验证集loss），当指标若干轮不再下降时自动降低学习率。这个调度器很适合不想手动调训练轮数的场景。需要注意，调度器的step调用时机不要放错：ReduceLROnPlateau的step需要传入监控指标，比如scheduler.step(val_loss)；而StepLR只需要scheduler.step()，通常放在每个epoch结束后调用。

4. Sequential、ModuleList与ModuleDict：三种模型组装方式，用了才知区别

4.1 nn.Sequential：适合搭建顺序流水线

nn.Sequential是初学者最先接触到的容器，把多个层按顺序串起来，前一个层的输出自动变成后一个层的输入。搭建一个简单MLP时，用Sequential非常直观：

import torch.nn as nn model = nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10) )

但Sequential的短板也明显：网络内部结构是线性的，不方便加跳跃连接、多分支结构。如果你要在两个层之间插入残差相加，或者让输入同时走两个不同分支再合并，Sequential就无能为力了。它最适合的是结构简单、顺序清晰的前馈网络。

4.2 nn.ModuleList：动态生成一堆同构层的神器

当你需要根据一个列表循环创建大量结构相同的层时，nn.ModuleList是最合适的选择。它像Python list一样支持append和索引，但和普通list不同——丢进去的子模块会被正确注册到网络中。我在构建深度可分离卷积的通道变换层时，就常常用ModuleList按通道数生成多个卷积层。

ModuleList有个使用细节：它本身不能直接做前向传播，你必须用循环遍历它，并手动把上一层输出传给下一个层。因为它只是管理子模块，不定义一个数据流顺序。顺序由你的forward代码控制，这给了你灵活性，代价是多写几行逻辑。

4.3 nn.ModuleDict：按名字选择层，动态路由更方便

ModuleDict按字典形式管理子模块。它可以让你通过一个字符串键来动态选择使用哪个层。在搭建多任务模型或需要根据配置切换到不同处理头的网络时，这种按名索引的方式特别方便。比如在不同的输入通道数量条件下，你可以用同一个模型类初始化不同配置的头部网络：

self.heads = nn.ModuleDict({ 'small': nn.Linear(64, 10), 'large': nn.Linear(256, 10) })

访问时self.heads["small"]就能拿到对应层。这个容器在配置驱动的模型构建场景中非常实用。

4.4 讲清楚一个底层机制：为什么普通list会丢参数

前面提过，nn.Module在__setattr__时会检查赋值的类型：如果值是nn.Module或nn.Parameter，就会自动放进_modules或_parameters字典。普通Python list不会被检查，所以里面的层根本不会注册到模型中。这也是为什么我建议所有把层放进容器的场景，都使用ModuleList或ModuleDict，而不是自己建list塞层。容器类本质上就是在帮你完成注册这个动作。

如果你实在需要自定义容器逻辑，也可以在自己写的类里手动调用self._modules['layer_name'] = layer来注册。但这属于高级操作，日常开发直接使用官方容器更稳妥。

5. 动手整合：用神经网络工具箱写一个能跑通的手写数字分类器

5.1 数据准备：从Dataset到DataLoader的完整链路

先准备好数据。以经典手写数字数据集为例，PyTorch把数据集下载、变换、分批加载的流程封装得非常顺滑。你需要定义或使用现成的Dataset类，返回一对数据：输入图像和对应标签。再用DataLoader按batch大小打乱并装载出批次数据：

from torch.utils.data import DataLoader from torchvision import datasets, transforms transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

Normalize的均值和标准差是数据集的全局统计量。一个常见的错误是拿这个归一化参数硬套到其他数据集上，导致输入分布不匹配。每换一个数据集，都应该重新统计对应的均值和标准差。

5.2 搭建CNN模型：卷积-池化-全连接的标准套路

用nn.Module写一个简单的CNN模型，正好把所有组件串起来：

import torch.nn.functional as F class SimpleCNN(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(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = self.fc2(x) return x

注意forward里最后没有softmax——前面讲CrossEntropyLoss时已经说明原因。x.view(x.size(0), -1)是把卷积特征图展平成向量，再送入全连接层。这里的展平操作是CNN转向全连接的关键，经常有新手忘记这步或者展平维度算错，导致矩阵乘法形状不匹配。

5.3 训练循环：model.train()和model.eval()的正确姿势

训练循环的核心结构其实很简单，但细节决定成败：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = SimpleCNN().to(device) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) criterion = nn.CrossEntropyLoss() for epoch in range(5): model.train() 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()

每次backward前都要zero_grad，否则梯度会跨batch累积，导致参数更新方向错乱。把模型和batch数据都放到GPU上也是必须的，如果模型在GPU而数据在CPU，前向传播时PyTorch会报设备不匹配的错误。

验证阶段就切换为model.eval()，并且用torch.inference_mode()包装，避免构建计算图：

model.eval() total_correct = 0 total = 0 with torch.inference_mode(): for images, labels in val_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) predicted = outputs.argmax(dim=1) total += labels.size(0) total_correct += (predicted == labels).sum().item()

这里argmax(dim=1)就是取每行最大logits对应的类别索引，作为预测结果。训练和验证两段代码分开写，逻辑清晰，也不容易串态。

5.4 模型保存与加载：规范动作，避免玄学问题

模型训练完毕，保存和加载几乎必然要用。我的标准做法是只保存state_dict：

torch.save(model.state_dict(), 'model_weights.pth') model = SimpleCNN() model.load_state_dict(torch.load('model_weights.pth')) model.eval()

很多人图省事直接torch.save(model, 'model.pth')，把整个模型对象序列化。这样做在模型类代码没变时能跑，但一旦类定义位置、导入路径、代码版本发生变化，加载时很容易出现兼容性错误。只保存state_dict的好处是模型结构在代码里定义，权重文件只负责存参数，解耦清晰。load_state_dict之后记得调用model.eval()，否则模型中的Dropout和BatchNorm还停留在训练模式，推理结果会不稳定。

6. 常见坑与排查思路：这些错误我都实际踩过

6.1 梯度完全不更新的排查链路

如果训练中loss一直原地不动，先别急着调学习率，按这个顺序排查：

• 检查是否调用了optimizer.zero_grad()，如果漏了，多个batch梯度叠加，更新会混乱。
• 检查loss.backward()是否真的执行了，尤其是用if条件控制梯度更新的代码路径时。
• 检查optimizer的参数列表是否来自model.parameters()，如果用错了模型，更新的一方和计算loss的一方参数对不上。
• 检查模型里的层是否都注册成了nn.Module，普通list这个问题不报错但权重不更新。

我在调试这类问题时，最常用的办法是打印模型参数在step前后的数值：

for name, param in model.named_parameters(): print(name, param.grad is not None)

如果grad始终为None，问题几乎都出在前向传播断了计算图的连接，比如把某个中间张量转成了numpy数字，或者用了不支持梯度计算的自定义操作。

6.2 显存OOM与累积梯度：一个容易被忽略的细节

显存不够是训练时的家常便饭。除了减小batch size，还有一个非常实用的技巧——梯度累积。由于显存限制，你没法一次塞进比较大的batch，于是可以把多个小batch的梯度累加起来再更新一次参数，等效于增大了batch size。实现方式是把zero_grad和step的频率拉低，每累积固定步数才更新一次：

accumulation_steps = 4 for idx, (images, labels) in enumerate(train_loader): loss = criterion(model(images), labels) / accumulation_steps loss.backward() if (idx + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()

这里把loss除以accumulation_steps再backward，是为了让累积后的梯度量级和直接用一个完整大batch训练时保持一致。曾经有人直接加起来不除，结果等效学习率变大好几倍，训练曲线直接飞了。

6.3 数据集加载慢与DataLoader参数调优

很多人用DataLoader时默认参数一挂就训练，结果发现GPU利用率一直上不去。影响数据加载速度的关键参数是num_workers和prefetch_factor。在Linux系统上把num_workers设成CPU核心数的一半以上，可以明显提升数据吞吐；Windows系统上，建议num_workers不要开太多，否则容易出现莫名其妙的报错。如果你的训练流程非常简单，可以设pin_memory=True，当使用GPU时，这会把数据锁定在页锁内存中，减少数据从CPU传到GPU的耗时。

对于超大型数据集，如果每次训练都要读几百GB的数据，还可以考虑用LMDB或HDF5格式做IO优化，这是后话。但对于大多数中小型项目，把num_workers和pin_memory调对，训练速度的提升就已经很明显了。

6.4 model.eval()遗漏：推理结果不稳定的常见源头

我在一开始做模型部署时多次吃过这个亏：训练时准确率很高，但单独跑推理脚本时，结果batch之间波动很大，尤其是用了BatchNorm和Dropout之后，每次推理输出都不一样。后来一排查，就是忘了在加载模型后调用model.eval()。神经网络工具箱里的这些模式切换，不是可有可无的API，而是直接决定模型计算行为的开关。把这个原则记住：训练模型用model.train()，推理和验证用model.eval()，刚加载的模型也要eval，这样才能保证输出稳定。

如果再往上走一步，torch.inference_mode()相比torch.no_grad()是更新的推荐方式，它进一步关闭了视图相关的追踪机制，速度更快。只要不涉及需要保留梯度的操作，我都建议用inference_mode，写推理代码时可以把它当成默认选择。

7. 写在最后：神经网络工具箱的正确打开方式

我自己的学习路径是：先照着官方文档把nn.Linear、nn.Conv2d、nn.CrossEntropyLoss、optim.Adam这些常用组件一个个用过一遍，然后强迫自己手写一个包含自定义层、动态容器、多分支结构的模型，再回头去读nn.Module的实现细节。经过这个过程，神经网络的很多“感觉”才真正建立起来。你在实际写代码时一定还会碰到新的坑，但只要掌握了排查的思路——先看数据流、再看梯度流、再看设备流——绝大多数问题都能很快定位。第3章这块内容就是把工具箱里的工具一件件拆开给你看，希望这篇笔记能让你少走点弯路，把更多时间花在真正有趣的模型设计上。