torch 中 nn.ModuleList()

nn.ModuleList() 是 PyTorch 中的一个类，用于管理神经网络模型中的子模块列表。它允许我们将多个子模块组织在一起，并将它们作为整个模型的一部分进行管理和操作。

在神经网络模型的开发过程中，通常需要定义和使用多个子模块，例如不同的层、块或者其他组件。nn.ModuleList() 提供了一种方便的方式来管理这些子模块，并确保它们被正确地注册为模型的一部分。

使用 nn.ModuleList() 需要进行两个步骤：

在模型的 __init__ 方法中，定义一个 nn.ModuleList 实例，并将需要管理的子模块添加到该列表中。
在模型的 forward 方法中，使用 nn.ModuleList 实例来访问和操作子模块。
下面是一个示例，演示如何使用 nn.ModuleList：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        self.module_list = nn.ModuleList([
            nn.Linear(10, 20),
            nn.ReLU(),
            nn.Linear(20, 10),
        ])

    def forward(self, x):
        for module in self.module_list:
            x = module(x)
        return x

model = MyModel()
input_tensor = torch.randn(32, 10)
output_tensor = model(input_tensor)

在这个示例中，我们定义了一个名为 MyModel 的自定义模型类。在该类的 __init__ 方法中，我们创建了一个 nn.ModuleList 实例 module_list，并添加了三个子模块：一个线性层（nn.Linear）、一个 ReLU 激活函数（nn.ReLU）和另一个线性层。这些子模块将作为整个模型的一部分。

在模型的 forward 方法中，我们通过迭代 module_list 中的子模块，依次将输入数据 x 传递给它们，并获取最终的输出。

通过使用 nn.ModuleList，我们可以方便地管理模型中的多个子模块，并确保它们被正确地注册为模型的一部分。这使得模型的结构清晰可见，同时也方便了模型的训练和参数优化。

torch 中 nn.ModuleList()

猜你喜欢