PyTorch:神经网络模块

PyTorch:神经网络模块 PyTorch 的 torch.nn 模块是构建神经网络最核心的模块。它提供了模型基类、常用网络层、激活函数、损失函数、容器结构、参数管理和训练状态控制等能力。简单地说torch.nn 模块回答的是神经网络在 PyTorch 中如何表示、如何组织、如何完成前向计算以及模型参数如何被自动管理。如果说 torch 模块负责张量与基础计算那么 torch.nn 模块就负责把这些张量计算组织成“模型”。在实际深度学习项目中我们很少只写零散的矩阵乘法而是会把线性层、卷积层、激活函数、归一化层、Dropout 层和损失函数组合成一个完整的神经网络。一、认识 torch.nn 模块torch.nn 是 PyTorch 中专门用于神经网络建模的模块。它并不是只提供几个网络层而是围绕模型构建提供了一套完整机制。图 1torch.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()一个最简单的使用方式如下import torchimport torch.nn as nn # 定义一个线性层输入特征数为 4输出特征数为 3layer nn.Linear(4, 3) # 构造 2 个样本每个样本有 4 个特征x torch.randn(2, 4) # 前向计算得到 2 个样本在 3 个输出维度上的结果y layer(x) print(y.shape)输出形状为torch.Size([2, 3])这里可以看出nn.Linear 不只是一个普通函数而是一个带有可学习参数的模块。它内部保存了权重和偏置并会在训练过程中被更新。二、nn.Module所有神经网络模块的基础在 PyTorch 中大多数神经网络模型都会继承 nn.Module。它是模型组织的基础。图 2nn.Module 的基本结构一个典型模型通常包含两个部分• __init__()定义模型中需要用到的层• forward()定义数据如何从输入流向输出例如import torchimport torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() # 定义第一层4 个输入特征映射到 16 个隐藏特征 self.fc1 nn.Linear(4, 16) # 定义激活函数 self.relu nn.ReLU() # 定义第二层16 个隐藏特征映射到 3 个输出类别 self.fc2 nn.Linear(16, 3) def forward(self, x): # 输入先经过第一层 x self.fc1(x) # 再经过非线性激活 x self.relu(x) # 最后得到输出结果 x self.fc2(x) return x model SimpleNet() x torch.randn(8, 4)y model(x) print(y.shape) # torch.Size([8, 3])这里最重要的不是代码长度而是理解模型的组织方式• __init__() 负责“有哪些层”• forward() 负责“数据怎么流动”• model(x) 会自动调用 forward(x)• 模型中的层会被自动注册参数也会被自动管理这就是 PyTorch 模型定义的基本范式。三、forward()定义前向传播逻辑forward() 方法决定输入张量如何一步步变成输出张量。它体现的是模型的计算路径。例如一个两层神经网络可以理解为输入 x→ 线性层 fc1→ 激活函数 ReLU→ 线性层 fc2→ 输出 logits对应代码是def forward(self, x): x self.fc1(x) x self.relu(x) x self.fc2(x) return x在分类任务中模型最后输出的通常不是直接类别而是一组分数常称为 logits。例如logits model(x)print(logits.shape)如果 logits.shape 是 (8, 3)可以理解为• 有 8 个样本• 每个样本对应 3 个类别分数• 分数越大模型越倾向于该类别需要注意的是很多分类模型的最后一层不需要手动加 Softmax。如果后面使用 nn.CrossEntropyLoss()它会在内部处理与分类概率相关的计算。因此模型通常直接输出原始分数即可。四、nn.Linear全连接层nn.Linear 是最基础、最常用的神经网络层之一。它常用于多层感知机、分类头、回归头和 Transformer 中的前馈网络部分。全连接层也称为“线性层”的核心计算可以理解为其中• X 表示输入张量• W 表示权重矩阵• b 表示偏置• Y 表示输出张量代码示例import torchimport torch.nn as nn # 输入特征数为 4输出特征数为 2linear nn.Linear(4, 2) # 5 个样本每个样本 4 个特征x torch.randn(5, 4) # 输出形状为 (5, 2)y linear(x) print(y.shape)print(linear.weight.shape)print(linear.bias.shape)需要注意的是nn.Linear(4, 2) 中的 4 和 2 分别表示• in_features4每个样本输入特征数• out_features2每个样本输出特征数它并不表示有 4 个样本或 2 个样本。样本数量通常由输入张量的第 0 维决定。五、激活函数引入非线性能力如果神经网络只有线性层那么多层线性变换叠加后本质上仍然可以合并成一个线性变换。为了让模型表达更复杂的非线性关系需要引入激活函数。常见激活函数包括• nn.ReLU• nn.Sigmoid• nn.Tanh• nn.LeakyReLU• nn.GELU• nn.Softmax示例import torchimport torch.nn as nn x torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]) # ReLU 会把负数变为 0正数保持不变relu nn.ReLU() # Sigmoid 会把数值压缩到 0 到 1 之间sigmoid nn.Sigmoid() # Tanh 会把数值压缩到 -1 到 1 之间tanh nn.Tanh() print(relu(x))print(sigmoid(x))print(tanh(x))在现代深度学习中ReLU 和它的变体非常常见在 Transformer、BERT、GPT 等模型中GELU 也经常出现。需要注意的是激活函数并不是随便添加的。它通常位于线性层或卷积层之后用于增加模型的非线性表达能力。六、nn.Sequential按顺序组织模型如果模型结构是简单的“从前到后依次执行”可以使用 nn.Sequential 简化代码。例如下面这个模型与前面手写 forward() 的两层网络等价import torchimport torch.nn as nn model nn.Sequential( nn.Linear(4, 16), # 输入层到隐藏层 nn.ReLU(), # 非线性激活 nn.Linear(16, 3) # 隐藏层到输出层) x torch.randn(8, 4)y model(x) print(y.shape)nn.Sequential 的优点是简洁适合线性堆叠结构。例如• 多层感知机• 简单卷积网络• 分类头• 回归头• 特征映射模块但如果模型中包含复杂分支、跳跃连接、多输入多输出或条件逻辑nn.Sequential 就不够灵活。这时更适合继承 nn.Module 并手写 forward()。七、ModuleList 与 ModuleDict保存多个子模块有些模型中层的数量可能比较多或者需要在循环中组织多个子模块。这时可以使用 nn.ModuleList 或 nn.ModuleDict。1、ModuleList列表形式保存子模块import torchimport torch.nn as nn class DeepMLP(nn.Module): def __init__(self): super().__init__() # 用 ModuleList 保存多个线性层 self.layers nn.ModuleList([ nn.Linear(4, 16), nn.Linear(16, 16), nn.Linear(16, 3) ]) self.relu nn.ReLU() def forward(self, x): # 前两层后接 ReLU最后一层直接输出 for layer in self.layers[:-1]: x self.relu(layer(x)) x self.layers[-1](x) return x model DeepMLP() x torch.randn(8, 4)y model(x) print(y.shape) # torch.Size([8, 3])ModuleList 的关键作用是让 PyTorch 能够正确识别并管理其中的子模块参数。不要简单地用普通 Python 列表保存网络层。普通列表中的层可能不会被模型正确注册从而导致参数无法被优化器更新。2、ModuleDict字典形式保存子模块import torchimport torch.nn as nn class MultiHeadModel(nn.Module): def __init__(self): super().__init__() self.shared nn.Linear(4, 16) # 用 ModuleDict 保存多个任务头 self.heads nn.ModuleDict({ classification: nn.Linear(16, 3), regression: nn.Linear(16, 1) }) self.relu nn.ReLU() def forward(self, x, task): x self.relu(self.shared(x)) x self.heads[task](x) return x model MultiHeadModel() x torch.randn(8, 4) y_cls model(x, taskclassification)y_reg model(x, taskregression) print(y_cls.shape) # torch.Size([8, 3])print(y_reg.shape) # torch.Size([8, 1])ModuleDict 适合需要按名称选择不同子模块的场景例如多任务学习、多个输出头、多种特征处理分支等。八、损失函数衡量预测与目标的差异神经网络训练不仅需要模型输出还需要损失函数。损失函数用于衡量模型预测结果与真实目标之间的差异。常见损失函数包括• nn.CrossEntropyLoss多分类任务常用• nn.BCELoss二分类概率输出场景• nn.BCEWithLogitsLoss二分类或多标签任务常用• nn.MSELoss回归任务常用• nn.L1Loss平均绝对误差损失• nn.NLLLoss负对数似然损失1、分类任务中的 CrossEntropyLossimport torchimport torch.nn as nn class DeepMLP(nn.Module): def __init__(self): super().__init__() # 用 ModuleList 保存多个线性层 self.layers nn.ModuleList([ nn.Linear(4, 16), nn.Linear(16, 16), nn.Linear(16, 3) ]) self.relu nn.ReLU() def forward(self, x): # 前两层后接 ReLU最后一层直接输出 for layer in self.layers[:-1]: x self.relu(layer(x)) x self.layers[-1](x) return x model DeepMLP() x torch.randn(8, 4)y model(x) print(y.shape) # torch.Size([8, 3])这里需要特别注意• logits 的形状通常是 (batch_size, num_classes)• target 的形状通常是 (batch_size,)• target 中保存的是类别编号不是 one-hot 向量• 使用 CrossEntropyLoss 时模型最后通常不需要手动加 Softmax这是初学者非常容易出错的地方。如果把标签写成 one-hot或者提前对 logits 做了不必要的 Softmax可能会导致训练效果异常。2、回归任务中的 MSELossimport torchimport torch.nn as nn # 8 个样本每个样本预测 1 个连续数值pred torch.randn(8, 1) # 对应的真实连续数值target torch.randn(8, 1) # 均方误差损失criterion nn.MSELoss() loss criterion(pred, target) print(loss)回归任务中预测值和真实值通常都使用浮点张量并且形状应尽量保持一致。例如如果 pred 的形状是 (8, 1)而 target 的形状是 (8,)有时可能会触发广播导致计算含义不符合预期。因此回归任务中要特别注意预测值和目标值的形状。九、卷积层、池化层与图像特征提取在图像任务中常用的神经网络层包括卷积层、归一化层、激活函数和池化层。典型组合可以写成import torchimport torch.nn as nn class SmallCNN(nn.Module): def __init__(self): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), # 输入 3 通道输出 16 通道 nn.BatchNorm2d(16), # 对 16 个通道做批量归一化 nn.ReLU(), # 引入非线性 nn.MaxPool2d(kernel_size2), # 高宽减半 nn.Conv2d(16, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(kernel_size2) ) self.classifier nn.Linear(32 * 8 * 8, 10) def forward(self, x): x self.features(x) # 保留批次维度把每张图像的特征展平成向量 x torch.flatten(x, start_dim1) x self.classifier(x) return x model SmallCNN() # 8 张 32×32 RGB 图像x torch.randn(8, 3, 32, 32) y model(x) print(y.shape) # torch.Size([8, 10])这里的结构可以理解为• 卷积层负责提取局部特征• 激活函数引入非线性• 池化层降低空间尺寸• 展平操作把特征图变成向量• 全连接层输出分类结果在卷积网络中输入通常采用 (N, C, H, W) 格式。如果图像数据是 (N, H, W, C)需要先调整维度顺序。十、Dropout 与归一化层在神经网络训练中除了线性层和卷积层还经常使用 Dropout 与归一化层来改善训练效果。1、Dropout随机失活import torchimport torch.nn as nn dropout nn.Dropout(p0.5) x torch.ones(10) # 训练模式下部分元素会被随机置为 0dropout.train()print(dropout(x)) # 推理模式下Dropout 不再随机丢弃元素dropout.eval()print(dropout(x))Dropout 常用于缓解过拟合。它在训练阶段随机丢弃部分神经元输出使模型不要过度依赖某些局部特征。需要注意的是Dropout 在训练模式和推理模式下行为不同。因此训练时要调用 model.train()推理或验证时要调用 model.eval()。2、BatchNorm 与 LayerNorm常见归一化层包括• nn.BatchNorm1d• nn.BatchNorm2d• nn.LayerNorm示例import torchimport torch.nn as nn # 常用于全连接网络或一维特征bn1d nn.BatchNorm1d(16) # 常用于卷积网络中的特征图bn2d nn.BatchNorm2d(16) # 常用于 Transformer 等结构ln nn.LayerNorm(16) x1 torch.randn(8, 16)x2 torch.randn(8, 16, 32, 32) print(bn1d(x1).shape)print(bn2d(x2).shape)归一化层的作用不是改变任务目标而是让中间特征的数值分布更稳定从而帮助模型更顺利地训练。十一、Embedding、RNN 与 Transformer 相关层torch.nn 不只用于普通前馈网络和卷积网络也提供了文本、序列和注意力模型相关的常用层。1、Embedding把离散编号映射为向量nn.Embedding 常用于自然语言处理和推荐系统。它可以把词编号、用户编号、物品编号等离散 ID 映射为连续向量。示例import torchimport torch.nn as nn # 词表大小为 1000每个词映射为 64 维向量embedding nn.Embedding(num_embeddings1000, embedding_dim64) # 4 个样本每个样本包含 10 个词编号token_ids torch.randint(0, 1000, (4, 10)) # 输出形状为 (4, 10, 64)x embedding(token_ids) print(x.shape)这里可以理解为• 输入是整数编号• 输出是连续向量• 每个编号都有一个可学习的向量表示2、RNN、LSTM 与 GRU处理序列数据的循环网络结构RNN、LSTM 与 GRU 都属于处理序列数据的循环神经网络结构常用于文本、语音、时间序列等任务。RNN循环神经网络是基础循环结构LSTM长短期记忆网络和 GRU门控循环单元是对普通 RNN 的改进版本目的都是让模型更好地处理序列中的上下文信息。示例import torchimport torch.nn as nn # 输入特征维度为 32隐藏状态维度为 64rnn nn.LSTM(input_size32, hidden_size64, batch_firstTrue) # 8 个样本每个样本长度为 10每个时间步 32 个特征x torch.randn(8, 10, 32) # output 保存每个时间步的输出h_n 和 c_n 保存最终隐藏状态与记忆状态output, (h_n, c_n) rnn(x) print(output.shape)print(h_n.shape)print(c_n.shape)当 batch_firstTrue 时输入形状通常是(batch_size, sequence_length, input_size)序列模型中最容易混淆的是维度顺序。写代码时要特别确认• 哪一维是批次• 哪一维是时间步• 哪一维是特征3、MultiheadAttention 与 Transformernn.MultiheadAttention 和 nn.Transformer 可用于构建注意力机制与 Transformer 结构。import torchimport torch.nn as nn # 嵌入维度为 64注意力头数为 4attention nn.MultiheadAttention(embed_dim64, num_heads4, batch_firstTrue) # 8 个样本序列长度为 10每个 token 是 64 维向量x torch.randn(8, 10, 64) # 自注意力中query、key、value 可以来自同一个输入output, weights attention(x, x, x) print(output.shape)print(weights.shape)这里重点是理解• 输入通常是序列张量• 每个位置会与其他位置建立关系• 输出形状通常与输入主结构保持一致Transformer 结构较复杂初学阶段不必急于展开所有细节。先理解 Embedding、序列张量形状和注意力输入输出关系会更稳妥。十二、参数管理parameters() 与 state_dict()神经网络训练的本质是不断更新模型中的可学习参数。nn.Module 会自动管理模型中的参数。1、查看模型参数import torch.nn as nn model nn.Sequential( nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3)) # 遍历模型中所有可学习参数for name, param in model.named_parameters(): print(name, param.shape)输出结果类似0.weight torch.Size([16, 4])0.bias torch.Size([16])2.weight torch.Size([3, 16])2.bias torch.Size([3])这里可以看到ReLU 没有可学习参数而 Linear 有权重和偏置。2、parameters()交给优化器更新import torch.optim as optim model nn.Linear(4, 3) # 优化器通过 model.parameters() 获取需要更新的参数optimizer optim.Adam(model.parameters(), lr0.001)这也是为什么模型层必须正确注册到 nn.Module 中。如果某些层没有被注册优化器可能无法找到它们的参数。3、state_dict()保存模型参数import torchimport torch.nn as nn model nn.Linear(4, 3) # 查看模型参数字典state model.state_dict() print(state.keys()) # 保存模型参数torch.save(model.state_dict(), model.pth)加载参数通常写成model nn.Linear(4, 3) # 加载已保存的参数model.load_state_dict(torch.load(model.pth))state_dict() 是 PyTorch 中非常常用的模型保存方式。它保存的是参数而不是完整训练代码。因此加载时通常需要先重新定义相同结构的模型。十三、train() 与 eval()训练模式和推理模式某些层在训练和推理时行为不同例如• Dropout训练时随机丢弃推理时不丢弃• BatchNorm训练时使用当前批次统计量推理时使用累计统计量因此在训练和验证时要明确切换模式。# 进入训练模式model.train() # 训练阶段通常会执行前向传播、计算损失、反向传播和参数更新logits model(x)loss criterion(logits, target) # 进入推理模式model.eval() # 推理阶段通常还会配合 torch.no_grad()with torch.no_grad(): logits model(x)需要注意• model.train() 不会自动训练模型只是切换为训练模式• model.eval() 不会停止梯度计算只是切换为推理模式• 推理时通常还要使用 torch.no_grad() 关闭梯度记录十四、一个完整示例用 nn 构建并训练简单分类模型图 3torch.nn 模型训练闭环下面用一个完整示例把 torch.nn、损失函数和优化器串起来。import torchimport torch.nn as nnimport torch.optim as optim # 1. 构造模拟数据200 个样本每个样本 4 个特征类别数为 3X torch.randn(200, 4)y torch.randint(0, 3, (200,)) class MLPClassifier(nn.Module): def __init__(self): super().__init__() # 使用 Sequential 组织简单多层感知机 self.net nn.Sequential( nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3) ) def forward(self, x): return self.net(x) # 2. 创建模型model MLPClassifier() # 3. 多分类任务常用交叉熵损失criterion nn.CrossEntropyLoss() # 4. Adam 优化器负责更新模型参数optimizer optim.Adam(model.parameters(), lr0.01) # 5. 训练模型for epoch in range(10): model.train() # 前向传播 logits model(X) # 计算损失 loss criterion(logits, y) # 清空上一轮梯度 optimizer.zero_grad() # 反向传播 loss.backward() # 更新参数 optimizer.step() print(f第 {epoch 1} 轮损失{loss.item():.4f})这个示例体现了神经网络训练的基本闭环输入数据 → 模型前向计算 → 计算损失 → 清空梯度 → 反向传播 → 参数更新其中torch.nn 主要负责模型结构和损失函数优化器负责根据梯度更新参数。十五、使用 torch.nn 时应注意的问题1、不要忘记继承 nn.Module 并调用 super().init()自定义模型时通常要写class MyModel(nn.Module): def __init__(self): super().__init__()如果忘记调用 super().__init__()子模块和参数可能无法被正常注册后续训练会出现问题。2、网络层应定义在 init() 中推荐写法class MyModel(nn.Module): def __init__(self): super().__init__() self.fc nn.Linear(4, 3) def forward(self, x): return self.fc(x)不推荐在 forward() 中反复新建带参数的层def forward(self, x): # 不推荐每次前向传播都会新建一层 fc nn.Linear(4, 3) return fc(x)原因是训练需要持续更新同一组参数。如果每次前向传播都创建新层模型参数就无法稳定学习。3、注意输入输出形状神经网络层对输入形状有明确要求。例如• nn.Linear 通常关注输入最后一维• nn.Conv2d 通常要求输入为 (N, C, H, W)• nn.LSTM 在 batch_firstTrue 时常用 (N, T, D)• nn.CrossEntropyLoss 常用 (N, C) 的预测值和 (N,) 的类别标签很多错误都可以通过打印形状定位print(x.shape)print(logits.shape)print(target.shape)4、区分 logits、概率和类别分类模型中常见三个概念• logits模型输出的原始分数• 概率经过 Softmax 后得到的归一化结果• 类别分数最大或概率最大的类别编号例如import torch logits torch.randn(4, 3) # 将原始分数转换为概率probs torch.softmax(logits, dim1) # 取概率最大的类别编号pred torch.argmax(probs, dim1) print(pred)训练时使用 CrossEntropyLoss通常直接输入 logits不需要先手动计算 Softmax。5、训练与推理模式要切换涉及 Dropout、BatchNorm 等层时训练和推理行为不同。训练阶段model.train()验证或推理阶段model.eval() with torch.no_grad(): output model(x)这两个操作不应省略。否则验证结果可能不稳定推理结果也可能受到训练行为影响。6、用 ModuleList 保存多个层而不是普通列表如果多个层需要放在列表里推荐使用self.layers nn.ModuleList([ nn.Linear(4, 16), nn.Linear(16, 3)])不要简单写成self.layers [ nn.Linear(4, 16), nn.Linear(16, 3)]普通列表中的层可能不会被正确注册优化器也可能无法更新其中的参数。7、损失函数要匹配任务类型不同任务应该使用不同损失函数• 多分类任务常用 CrossEntropyLoss• 二分类任务常用 BCEWithLogitsLoss• 回归任务常用 MSELoss 或 L1Loss• 多标签任务常用 BCEWithLogitsLoss损失函数选错模型即使能运行也可能学不到正确目标。 小结torch.nn 模块负责神经网络模型的构建与组织是从张量计算走向深度学习建模的关键模块。学习这一模块重点不只是记住层名称而是理解 nn.Module 如何管理层和参数forward() 如何定义数据流损失函数如何连接预测与目标以及训练模式和推理模式为何必须区分。“点赞有美意赞赏是鼓励”