torch.nn.ModuleList 是 PyTorch 中的一个容器模块，用于存储子模块（torch.nn.Module 对象）的列表。与 Python 的普通列表不同，ModuleList 是一个特殊的容器，它能够正确地注册其包含的子模块，使得这些子模块的参数可以被 PyTorch 的优化器识别和更新。

主要特点

  1. 自动注册子模块：ModuleList 会自动将其包含的所有子模块注册到父模块中。这意味着这些子模块的参数会被 PyTorch 的优化器识别，并且在调用 model.parameters() 时会被包含在内。
  2. 动态添加子模块：你可以在 ModuleList 中动态地添加或删除子模块，而不需要重新定义整个模型。
  3. 索引和迭代：ModuleList 支持类似于 Python 列表的索引和迭代操作，因此你可以方便地访问其中的子模块。

使用场景

ModuleList 通常用于以下场景：

  • 动态网络结构：当你需要在模型中动态地添加或删除层时，ModuleList 是一个很好的选择。
  • 重复的子模块：如果你有多个相似的子模块（例如多个卷积层），你可以将它们存储在 ModuleList 中，并通过循环来访问它们。

示例代码

以下是一个简单的示例，展示了如何使用 ModuleList：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # 使用 ModuleList 存储多个线性层
        self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(5)])
    
    def forward(self, x):
        # 依次通过所有的线性层
        for layer in self.layers:
            x = layer(x)
        return x

# 创建模型实例
model = MyModel()

# 打印模型结构
print(model)

# 打印模型参数
for param in model.parameters():
    print(param)

注意事项

  • 不能直接用于计算：ModuleList 本身并不参与前向传播计算，它只是一个容器。你需要手动在 forward 方法中调用其中的子模块。
  • 与 Sequential 的区别：nn.Sequential 是一个顺序容器，它会自动将输入数据依次通过所有的子模块。而 ModuleList 只是一个简单的列表容器，不会自动进行前向传播。

总结

torch.nn.ModuleList 是一个非常有用的工具，特别是在你需要动态管理模型中的子模块时。它能够确保子模块的参数被正确注册，并且提供了类似于 Python 列表的灵活操作方式。