PyTorch：神经网络模块

PyTorch 的 torch.nn 模块，是构建神经网络最核心的模块。它提供了模型基类、常用网络层、激活函数、损失函数、容器结构、参数管理和训练状态控制等能力。

简单地说，torch.nn 模块回答的是：神经网络在 PyTorch 中如何表示、如何组织、如何完成前向计算，以及模型参数如何被自动管理。

如果说 torch 模块负责张量与基础计算，那么 torch.nn 模块就负责把这些张量计算组织成“模型”。在实际深度学习项目中，我们很少只写零散的矩阵乘法，而是会把线性层、卷积层、激活函数、归一化层、Dropout 层和损失函数组合成一个完整的神经网络。

一、认识 torch.nn 模块

torch.nn 是 PyTorch 中专门用于神经网络建模的模块。它并不是只提供几个网络层，而是围绕模型构建提供了一套完整机制。

图 1：torch.nn 模块在 PyTorch 训练流程中的位置

在 PyTorch 中，一个神经网络通常由以下部分组成：

• 模型基类：nn.Module

• 网络层：如 nn.Linear、nn.Conv2d、nn.Embedding

• 激活函数：如 nn.ReLU、nn.Sigmoid、nn.GELU

• 损失函数：如 nn.CrossEntropyLoss、nn.MSELoss

• 容器结构：如 nn.Sequential、nn.ModuleList

• 参数管理：如 parameters()、state_dict()

• 训练与推理状态：如 train()、eval()

一个最简单的使用方式如下：

输出形状为：

torch.Size([2, 3])

这里可以看出，nn.Linear 不只是一个普通函数，而是一个带有可学习参数的模块。它内部保存了权重和偏置，并会在训练过程中被更新。

二、nn.Module：所有神经网络模块的基础

在 PyTorch 中，大多数神经网络模型都会继承 nn.Module。它是模型组织的基础。

图 2：nn.Module 的基本结构

一个典型模型通常包含两个部分：

• __init__()：定义模型中需要用到的层

• forward()：定义数据如何从输入流向输出

例如：

这里最重要的不是代码长度，而是理解模型的组织方式：

• __init__() 负责“有哪些层”

• forward() 负责“数据怎么流动”

• model(x) 会自动调用 forward(x)

• 模型中的层会被自动注册，参数也会被自动管理

这就是 PyTorch 模型定义的基本范式。

三、forward()：定义前向传播逻辑

forward() 方法决定输入张量如何一步步变成输出张量。它体现的是模型的计算路径。

例如，一个两层神经网络可以理解为：

→ 输出 logits

对应代码是：

在分类任务中，模型最后输出的通常不是直接类别，而是一组分数，常称为 。

例如：

如果 logits.shape 是 (8, 3)，可以理解为：

• 有 8 个样本

• 每个样本对应 3 个类别分数

• 分数越大，模型越倾向于该类别

需要注意的是，很多分类模型的最后一层不需要手动加 。如果后面使用 nn.CrossEntropyLoss()，它会在内部处理与分类概率相关的计算。因此，模型通常直接输出原始分数即可。

四、nn.Linear：全连接层

nn.Linear 是最基础、最常用的神经网络层之一。它常用于多层感知机、分类头、回归头和 Transformer 中的前馈网络部分。

（也称为“线性层”）的核心计算可以理解为：

其中：

• X 表示输入张量

• W 表示权重矩阵

• b 表示偏置

• Y 表示输出张量

代码示例：

需要注意的是，nn.Linear(4, 2) 中的 4 和 2 分别表示：

• in_features=4：每个样本输入特征数

• out_features=2：每个样本输出特征数

它并不表示有 4 个样本或 2 个样本。样本数量通常由输入张量的第 0 维决定。

五、激活函数：引入非线性能力

如果神经网络只有线性层，那么多层线性变换叠加后，本质上仍然可以合并成一个线性变换。为了让模型表达更复杂的非线性关系，需要引入。

常见激活函数包括：

• nn.ReLU

• nn.Sigmoid

• nn.Tanh

• nn.LeakyReLU

• nn.GELU

• nn.Softmax

示例：

在现代深度学习中， 和它的变体非常常见；在 Transformer、BERT、GPT 等模型中， 也经常出现。

需要注意的是，激活函数并不是随便添加的。它通常位于线性层或卷积层之后，用于增加模型的非线性表达能力。

六、nn.Sequential：按顺序组织模型

如果模型结构是简单的“从前到后依次执行”，可以使用 nn.Sequential 简化代码。

例如，下面这个模型与前面手写 forward() 的两层网络等价：

nn.Sequential 的优点是简洁，适合线性堆叠结构。例如：

• 多层感知机

• 简单卷积网络

• 分类头

• 回归头

• 特征映射模块

但如果模型中包含复杂分支、跳跃连接、多输入多输出或条件逻辑，nn.Sequential 就不够灵活。这时更适合继承 nn.Module 并手写 forward()。

七、ModuleList 与 ModuleDict：保存多个子模块

有些模型中，层的数量可能比较多，或者需要在循环中组织多个子模块。这时可以使用 nn.ModuleList 或 nn.ModuleDict。

1、ModuleList：列表形式保存子模块

ModuleList 的关键作用是：让 PyTorch 能够正确识别并管理其中的子模块参数。

不要简单地用普通 Python 列表保存网络层。普通列表中的层可能不会被模型正确注册，从而导致参数无法被优化器更新。

2、ModuleDict：字典形式保存子模块

ModuleDict 适合需要按名称选择不同子模块的场景，例如多任务学习、多个输出头、多种特征处理分支等。

八、损失函数：衡量预测与目标的差异

神经网络训练不仅需要模型输出，还需要。损失函数用于衡量模型预测结果与真实目标之间的差异。

常见损失函数包括：

• nn.CrossEntropyLoss：多分类任务常用

• nn.BCELoss：二分类概率输出场景

• nn.BCEWithLogitsLoss：二分类或多标签任务常用

• nn.MSELoss：回归任务常用

• nn.L1Loss：平均绝对误差损失

• nn.NLLLoss：负对数似然损失

1、分类任务中的 CrossEntropyLoss

nn.CrossEntropyLoss 常用于多分类任务。它接收的是模型输出的原始分数 logits，以及真实类别编号 target。

这里需要特别注意：

• logits 的形状通常是 (batch_size, num_classes)

• target 的形状通常是 (batch_size,)

• target 中保存的是类别编号，而不是 one-hot 向量；target 应是整数类型张量，常见类型是 torch.long

需要注意的是，CrossEntropyLoss 内部已经包含了与 Softmax 相关的计算，因此模型最后一层通常直接输出 logits，不需要手动添加 Softmax。

下面这种写法更符合 PyTorch 多分类任务的常见形式：

这是初学者非常容易出错的地方。如果提前对 logits 做 Softmax，或者把 target 写成 one-hot 向量，就可能导致损失计算语义错误，进而影响训练效果。

2、回归任务中的 MSELoss

回归任务中，预测值和真实值通常都使用浮点张量，并且形状应尽量保持一致。

例如，如果 pred 的形状是 (8, 1)，而 target 的形状是 (8,)，有时可能会触发广播，导致计算含义不符合预期。因此，回归任务中要特别注意预测值和目标值的形状。

九、卷积层、池化层与图像特征提取

在图像任务中，常用的神经网络层包括、、和。

典型组合可以写成：

这里的结构可以理解为：

• 卷积层负责提取局部特征

• 激活函数引入非线性

• 池化层降低空间尺寸

• 展平操作把特征图变成向量

• 全连接层输出分类结果

在卷积网络中，输入通常采用 (N, C, H, W) 格式。如果图像数据是 (N, H, W, C)，需要先调整维度顺序。

十、Dropout 与归一化层

在神经网络训练中，除了线性层和卷积层，还经常使用 与来改善训练效果。

1、Dropout：随机失活

Dropout 常用于缓解过拟合。它在训练阶段随机丢弃部分神经元输出，使模型不要过度依赖某些局部特征。

需要注意的是，Dropout 在训练模式和推理模式下行为不同。因此，训练时要调用 model.train()，推理或验证时要调用 model.eval()。

2、BatchNorm 与 LayerNorm

常见归一化层包括：

• nn.BatchNorm1d

• nn.BatchNorm2d

• nn.LayerNorm

示例：

归一化层的作用不是改变任务目标，而是让中间特征的数值分布更稳定，从而帮助模型更顺利地训练。

十一、Embedding、RNN 与 Transformer 相关层

torch.nn 不只用于普通前馈网络和卷积网络，也提供了文本、序列和注意力模型相关的常用层。

1、Embedding：把离散编号映射为向量

nn.Embedding 常用于自然语言处理和推荐系统。它可以把词编号、用户编号、物品编号等离散 ID 映射为连续向量。

示例：

这里可以理解为：

• 输入是整数编号

• 输出是连续向量

• 每个编号都有一个可学习的向量表示

2、RNN、LSTM 与 GRU：处理序列数据的循环网络结构

RNN、LSTM 与 GRU 都属于处理序列数据的循环神经网络结构，常用于文本、语音、时间序列等任务。

RNN（循环神经网络）是基础循环结构，（长短期记忆网络）和 （门控循环单元）是对普通 RNN 的改进版本，目的都是让模型更好地处理序列中的上下文信息。

示例：

当 batch_first=True 时，输入形状通常是：

(batch_size, sequence_length, input_size)

序列模型中最容易混淆的是维度顺序。写代码时要特别确认：

• 哪一维是批次

• 哪一维是时间步

• 哪一维是特征

3、MultiheadAttention 与 Transformer

nn.MultiheadAttention 和 nn.Transformer 可用于构建注意力机制与 Transformer 结构。

这里重点是理解：

• 输入通常是序列张量

• 每个位置会与其他位置建立关系

• 输出形状通常与输入主结构保持一致

Transformer 结构较复杂，初学阶段不必急于展开所有细节。先理解 Embedding、序列张量形状和注意力输入输出关系，会更稳妥。

十二、参数管理：parameters() 与 state_dict()

神经网络训练的本质，是不断更新模型中的可学习参数。nn.Module 会自动管理模型中的参数。

1、查看模型参数

输出结果类似：

2.bias torch.Size([3])

这里可以看到，ReLU 没有可学习参数，而 Linear 有权重和偏置。

2、parameters()：交给优化器更新

这也是为什么模型层必须正确注册到 nn.Module 中。如果某些层没有被注册，优化器可能无法找到它们的参数。

3、state_dict()：保存模型参数

加载参数通常写成：

state_dict() 是 PyTorch 中非常常用的模型保存方式。它保存的是参数，而不是完整训练代码。因此，加载时通常需要先重新定义相同结构的模型。

十三、train() 与 eval()：训练模式和推理模式

某些层在训练和推理时行为不同，例如：

• Dropout：训练时随机丢弃，推理时不丢弃

• BatchNorm：训练时使用当前批次统计量，推理时使用累计统计量

因此，在训练和验证时要明确切换模式。

需要注意：

• model.train() 不会自动训练模型，只是切换为训练模式

• model.eval() 不会停止梯度计算，只是切换为推理模式

• 推理时通常还要使用 torch.no_grad() 关闭梯度记录

十四、一个完整示例：用 nn 构建并训练简单分类模型

图 3：torch.nn 模型训练闭环

下面用一个完整示例，把 torch.nn、损失函数和优化器串起来。

这个示例体现了神经网络训练的基本闭环：

输入数据 → 模型前向计算 → 计算损失 → 清空梯度 → 反向传播 → 参数更新

其中，torch.nn 主要负责模型结构和损失函数；优化器负责根据梯度更新参数。

十五、使用 torch.nn 时应注意的问题

1、不要忘记继承 nn.Module 并调用 super().__init__()

自定义模型时，通常要写：

如果忘记调用 super().__init__()，子模块和参数可能无法被正常注册，后续训练会出现问题。

2、网络层应定义在 init() 中

推荐写法：

不推荐在 forward() 中反复新建带参数的层：

原因是：训练需要持续更新同一组参数。如果每次前向传播都创建新层，模型参数就无法稳定学习。

3、注意输入输出形状

神经网络层对输入形状有明确要求。例如：

• nn.Linear 通常关注输入最后一维

• nn.Conv2d 通常要求输入为 (N, C, H, W)

• nn.LSTM 在 batch_first=True 时常用 (N, T, D)

• nn.CrossEntropyLoss 常用 (N, C) 的预测值和 (N,) 的类别标签

很多错误都可以通过打印形状定位：

print(target.shape)

4、区分 logits、概率和类别

分类模型中，常见三个概念：

• logits：模型输出的原始分数

• 概率：经过 Softmax 后得到的归一化结果

• 类别：分数最大或概率最大的类别编号

例如：

训练时使用 CrossEntropyLoss，通常直接输入 logits，不需要先手动计算 Softmax。

5、训练与推理模式要切换

涉及 Dropout、BatchNorm 等层时，训练和推理行为不同。

训练阶段：

model.train()

验证或推理阶段：

    output = model(x)

这两个操作不应省略。否则，验证结果可能不稳定，推理结果也可能受到训练行为影响。

6、用 ModuleList 保存多个层，而不是普通列表

如果多个层需要放在列表里，推荐使用：

不要简单写成：

普通列表中的层可能不会被正确注册，优化器也可能无法更新其中的参数。

7、损失函数要匹配任务类型

不同任务应该使用不同损失函数：

• 多分类任务：常用 CrossEntropyLoss

• 二分类任务：常用 BCEWithLogitsLoss

• 回归任务：常用 MSELoss 或 L1Loss

• 多标签任务：常用 BCEWithLogitsLoss

损失函数选错，模型即使能运行，也可能学不到正确目标。

小结

torch.nn 模块负责神经网络模型的构建与组织，是从张量计算走向深度学习建模的关键模块。学习这一模块，重点不只是记住层名称，而是理解 nn.Module 如何管理层和参数，forward() 如何定义数据流，损失函数如何连接预测与目标，以及训练模式和推理模式为何必须区分。

“点赞有美意，赞赏是鼓励”