How to register a module in a list to be a submodule of a huge module?

Hello, I’m trying to create such a Network:

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.linears = [torch.nn.Linear(5, 10)] * 10
        self.special_linear = torch.nn.Linear(100, 500)

model = Net()

Because I need lots of linears so I just put them into a list at once, but it seems not be a submodule member of Net:

print(list(model.modules()))

/home/hdl2/anaconda3/bin/python /home/hdl2/Desktop/xxx/playground.py
[Net (
  (special_linear): Linear (100 -> 500)
), Linear (100 -> 500)]

Process finished with exit code 0

So I’d like to know if there any way to register those linears to be submodule member of Net? Thanks in advance.

        for i, l in enumerate(self.linears):
            super(Net, self).add_module(str(i) + 's', l)
1 Like

Alternatively, you could use nn.ModuleList, which acts like a list, but registers your modules:

self.linears = nn.ModuleList([nn.Linear(10, 10) for _ in range(10)])
12 Likes