Python PyTorch ModuleList用法及代碼示例

本文簡要介紹python語言中 torch.nn.ModuleList 的用法。

用法:

class torch.nn.ModuleList(modules=None)

參數：

modules(可迭代的,可選的) -要添加的可迭代模塊

在列表中保存子模塊。

ModuleList 可以像常規 Python 列表一樣進行索引，但它包含的模塊已正確注冊，並且所有 Module 方法都可見。

例子：

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])

    def forward(self, x):
        # ModuleList can act as an iterable, or be indexed using ints
        for i, l in enumerate(self.linears):
            x = self.linears[i // 2](x) + l(x)
        return x