Create dictionary of module layers

#1

In pytorch, we can create list of modules using ModuleList. Can we create a dictionary of modules? I need to set a name along with a layer. Any help would be appreciated.

#2

Have a look at the Sequential example. There they use an OrderedDict for this purpose.

Hope this helps!

#3

In that example, they wrapped the orderdict using nn.Sequential to make a single layer. In my case, I want to store a certain number of layers in a dictionary and later I want to use them based on conditions.

#4

The modules are stored as an ordered dictionary in nn.Sequential. You can directly access them with ._modules.

model = nn.Sequential(OrderedDict([
            ("layer_one", nn.Linear(1, 10)), ("layer_two", nn.Linear(10, 3))
        ]))
model._modules["layer_one"]
model._modules["layer_two"]
4 Likes
#5

We have ModuleDict now, which is easier to use :smile: