PyTorch中nn.Module详解

直接print(dir(nn.Module))，得到如下内容：

一、模型结构与参数

1. parameters()

   • 用途：返回模块的所有可训练参数（如权重、偏置）。
   • 示例：

     for param in model.parameters():
         print(param.shape)
     

2. named_parameters()

   • 用途：返回带名称的参数迭代器，便于调试和访问特定参数。
   • 示例：

     for name, param in model.named_parameters():
         if 'weight' in name:
             print(name, param.shape)
     

3. children()

   • 用途：返回直接子模块的迭代器。
   • 示例：

     for child in model.children():
         print(type(child))
     

4. modules()

   • 用途：递归返回所有子模块（包括自身）。
   • 示例：

     for module in model.modules():
         if isinstance(module, nn.Conv2d):
             print(module.kernel_size)
     

二、模型状态与模式

5. train() 和 eval()

   • 用途：切换训练/推理模式（影响Dropout、BatchNorm等层）。
   • 示例：

     model.train()  # 训练模式
     model.eval()   # 推理模式
     

6. training

   • 用途：布尔属性，指示当前模式（True 为训练，False 为推理）。
   • 示例：

     print(model.training)  # 输出：True/False
     

三、模型保存与加载

7. state_dict()

   • 用途：返回包含模型所有参数的字典（OrderedDict）。
   • 示例：

     torch.save(model.state_dict(), 'model.pth')
     

8. load_state_dict()

   • 用途：从字典加载模型参数。
   • 示例：

     model.load_state_dict(torch.load('model.pth'))
     

四、设备与数据类型

9. to()

   • 用途：将模型移动到指定设备（如GPU）或转换数据类型。
   • 示例：

     model.to('cuda')          # 移动到GPU
     model.to(torch.float16)   # 转换为半精度
     

10. cpu() 和 cuda()

    • 用途：快捷方法，分别将模型移动到CPU或GPU。
    • 示例：

      model.cuda()  # 等价于 model.to('cuda')
      

五、前向传播与计算

11. forward()

    • 用途：定义模型的前向传播逻辑（需在自定义模块中重写）。
    • 示例：

      class MyModel(nn.Module):
          def forward(self, x):
              return self.layer(x)
      

12. __call__()

    • 用途：调用模型实例时触发（内部调用 forward()，支持钩子函数）。
    • 示例：

      output = model(x)  # 等价于 output = model.forward(x)
      

六、参数初始化与优化

13. zero_grad()

    • 用途：清空所有参数的梯度（通常在每个训练步骤前调用）。
    • 示例：

      optimizer.zero_grad()  # 等价于 model.zero_grad()
      

14. requires_grad_()

    • 用途：设置参数是否需要梯度（用于冻结部分模型）。
    • 示例：

      for param in model.parameters():
          param.requires_grad = False  # 冻结所有参数
      

七、调试与信息

15. extra_repr()

    • 用途：自定义模块打印信息（需在子类中重写）。
    • 示例：

      class MyModel(nn.Module):
          def extra_repr(self):
              return f"hidden_size={self.hidden_size}"
      

16. dump_patches()

    • 用途：打印模型的补丁信息（用于调试版本差异）。

八、其他实用方法

17. apply()

    • 用途：递归应用函数到所有子模块（如初始化权重）。
    • 示例：

      def init_weights(m):
          if isinstance(m, nn.Conv2d):
              nn.init.kaiming_normal_(m.weight)
      model.apply(init_weights)
      

18. register_forward_hook()

    • 用途：注册前向传播钩子（用于捕获中间输出，调试或特征提取）。

总结

日常使用中，最频繁的方法包括：

• 模型构建：parameters(), children(), modules()
• 训练与推理：train(), eval(), zero_grad(), forward()
• 保存与加载：state_dict(), load_state_dict()
• 设备管理：to(), cuda(), cpu()

其他方法根据具体需求选择使用，例如钩子函数用于高级调试，apply() 用于统一初始化。

与nn.Sequential对比：

1. 继承关系与基础属性

• nn.Module

  • 是所有神经网络模块的基类，提供最基础的功能（如参数管理、钩子机制）。
  • 包含核心属性：_parameters, _modules, _buffers 等。

• nn.Sequential

  • 是 nn.Module 的子类，继承了所有基础功能。
  • 额外添加了与顺序执行相关的属性（如 __getitem__、append）。

2. 核心差异对比

3. 具体方法对比

3.1 公共方法（两者都有）

# 模型参数与结构
['parameters', 'named_parameters', 'children', 'modules', 'named_children', 'named_modules']

# 模型状态
['train', 'eval', 'training', 'zero_grad', 'requires_grad_']

# 设备与数据类型
['to', 'cpu', 'cuda', 'float', 'double', 'half', 'bfloat16']

# 保存与加载
['state_dict', 'load_state_dict']

# 钩子机制
['register_forward_hook', 'register_backward_hook']

3.2 nn.Sequential 特有的方法

# 列表操作（动态修改模块顺序）
['__getitem__', '__setitem__', '__delitem__', '__len__', 'append', 'extend', 'insert', 'pop']

# 索引相关
['_get_item_by_idx']

3.3 nn.Module 特有的方法

# 自定义实现
['forward', 'extra_repr']

# 高级管理
['add_module', 'register_module', 'register_parameter', 'register_buffer']

4. 示例对比

4.1 创建模型

# nn.Module（需自定义 forward）
class CustomModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        return self.relu(self.conv(x))

# nn.Sequential（自动按顺序执行）
seq_model = nn.Sequential(
    nn.Conv2d(3, 64, 3),
    nn.ReLU()
)

4.2 访问子模块

# nn.Module
custom_model.conv  # 通过属性名访问

# nn.Sequential
seq_model[0]       # 通过索引访问
seq_model.append(nn.MaxPool2d(2))  # 动态添加模块

5. 总结

建议：

• 对于简单的顺序网络，优先使用 nn.Sequential 以减少代码量。
• 对于包含复杂逻辑（如残差连接、多输入输出）的网络，使用 nn.Module 自定义实现。