介绍
nn.Module 是 PyTorch 中所有神经网络模块的 基类，自定义的任何神经网络层、整个模型，都需要 继承这个类。
它提供了一系列核心功能，无需手动处理参数管理、设备迁移、前向传播、模型保存/加载等繁琐工作。
简单来说：nn.Module 就像一个“神经网络工具箱”，封装了构建模型的所有基础能力。
功能与使用方法
基础使用步骤
自定义模型的标准流程是：
- 继承
nn.Module - 在
__init__中定义网络层（如卷积、全连接层） - 实现
forward方法（定义前向传播逻辑）
例子
import torch
import torch.nn as nn
# 自定义一个简单的全连接网络，继承 nn.Module
class MyNet(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
# 必须先调用父类的 __init__ 方法
super(MyNet, self).__init__()
# 定义网络层（nn.Module 会自动管理这些层的参数）
self.fc1 = nn.Linear(input_size, hidden_size) # 全连接层1
self.relu = nn.ReLU() # 激活函数
self.fc2 = nn.Linear(hidden_size, output_size) # 全连接层2
# 必须实现 forward 方法，定义数据的前向传播路径
def forward(self, x):
out = self.fc1(x) # 第一层线性变换
out = self.relu(out)# ReLU激活
out = self.fc2(out) # 第二层线性变换
return out
# ---------------------- 测试使用 ----------------------
# 1. 创建模型实例
model = MyNet(input_size=10, hidden_size=20, output_size=5)
print("模型结构：")
print(model)
# 2. 前向传播（直接调用模型实例，无需手动调用 forward 方法）
input_data = torch.randn(3, 10) # 3个样本，每个样本10维特征
output = model(input_data)
print("\n输入形状：", input_data.shape)
print("输出形状：", output.shape)
# 3. 查看模型参数（nn.Module 自动管理参数）
print("\n模型参数：")
for name, param in model.named_parameters():
print(f"参数名：{name}，形状：{param.shape}")
# 4. 设备迁移（nn.Module 一键移到 GPU/CPU）
if torch.cuda.is_available():
model = model.cuda()
input_data = input_data.cuda()
output = model(input_data)
print("\n迁移到 GPU 后的输出设备：", output.device)
关键功能解释
| 功能 | 说明 |
|---|---|
__init__ 继承 |
必须通过 super().__init__() 初始化父类，否则无法使用 nn.Module 的所有功能 |
forward 方法 |
定义数据的前向传播逻辑，调用模型时（如 model(x)）会自动执行该方法（无需手动调 forward） |
| 参数自动管理 | 所有在 __init__ 中定义的 nn.Module 子类（如 nn.Linear、nn.Conv2d）的参数，都会被 parameters()/named_parameters() 收集，用于后续的梯度下降 |
| 设备迁移 | 通过 model.to(device) 可以一键将模型所有参数移到 CPU/GPU，无需逐个处理层 |
| 模型保存/加载 | torch.save(model.state_dict(), "model.pth") 保存参数，model.load_state_dict(torch.load("model.pth")) 加载参数，这是 nn.Module 推荐的保存方式 |
| 子模块管理 | 可以通过 nn.Sequential 或直接定义子模块（如 self.feature = nn.Conv2d(3, 16, 3)），nn.Module 会自动管理所有子模块 |
常见误区
- 不要在
forward外定义层：如果层定义在forward里，nn.Module无法收集其参数，会导致参数无法更新； - 不要手动调用
forward：直接model(x)即可，PyTorch 会在调用时自动触发forward，并处理梯度计算等底层逻辑； - 区分
nn.Module和nn.functional：nn.Module（如nn.ReLU）是类，自带参数管理；nn.functional（如F.relu）是函数，无参数，需手动管理（适合无参数的激活函数/池化）。
总结
nn.Module是 PyTorch 构建神经网络的基类，所有模型/层都必须继承它；- 核心用法是：继承
nn.Module→ 在__init__定义层 → 实现forward定义前向传播； - 它提供了参数管理、设备迁移、模型保存/加载等核心功能，极大简化了神经网络开发。
掌握 nn.Module 是学好 PyTorch 的关键，所有复杂模型（如 ResNet、Transformer）本质上都是基于 nn.Module 构建的。