nn.ModuleDict vs OrderedDict vs dict

What is the advantage of using nn.ModuleDict versus OrderedDict or dict (in case the order doesn’t matter)?

self.activations = nn.ModuleDict([
        ['lrelu', nn.LeakyReLU()],
        ['prelu', nn.PReLU()]
])

vs

self.activations = OrderedDict([
        ['lrelu', nn.LeakyReLU()],
        ['prelu', nn.PReLU()]
])

nn.ModuleDict will properly register the modules to the parent as seen here:

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.module_activations = nn.ModuleDict([
            ['m_lrelu', nn.LeakyReLU()],
            ['m_prelu', nn.PReLU()]
        ])



        self.activations = OrderedDict([
            ['lrelu', nn.LeakyReLU()],
            ['prelu', nn.PReLU()]
        ])
    
model = MyModel()
print(dict(model.named_modules()))

This makes sure that the parameters are pushed to the desired device in e.g. to() calls and the model.parameters() call (used to pass the parameters to e.g. an optimizer) will return all parameters of the registered modules.

1 Like

The parameters may not be updated while training if you use OrderedDict.