In this tutorial, we will be discussing two important container classes in Pytorch: nn.ModuleList and nn.ModuleDict. These container classes are used to manage a list or a dictionary of nn.Module objects in a Pytorch model, making it easier to handle multiple layers or components.

nn.ModuleList:

nn.ModuleList is a container class that holds a list of nn.Module objects. It behaves like a regular Python list, but with some added functionality for Pytorch modules. You can add nn.Module objects to an nn.ModuleList using the append method, just like you would with a regular list.

Let’s see how to use nn.ModuleList in a neural network model:

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layers = nn.ModuleList([
            nn.Linear(10, 20),
            nn.ReLU(),
            nn.Linear(20, 1)
        ])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

model = MyModel()

In this example, we have defined a simple neural network model with three layers using nn.ModuleList. We have a linear layer, a ReLU activation function, and another linear layer. In the forward method, we iterate over each layer in the nn.ModuleList and apply them sequentially to the input x.

nn.ModuleDict:

nn.ModuleDict is a container class that holds a dictionary of nn.Module objects. It allows you to access the modules by their keys, making it convenient to manage multiple modules in a Pytorch model. You can add nn.Module objects to an nn.ModuleDict using dictionary-style indexing.

Let’s see how to use nn.ModuleDict in a neural network model:

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layers = nn.ModuleDict({
            'linear1': nn.Linear(10, 20),
            'relu': nn.ReLU(),
            'linear2': nn.Linear(20, 1)
        })

    def forward(self, x):
        x = self.layers['linear1'](x)
        x = self.layers['relu'](x)
        x = self.layers['linear2'](x)
        return x

model = MyModel()

In this example, we have defined a neural network model with three layers using nn.ModuleDict. We have defined each layer with a key in the dictionary, allowing us to access the modules by their keys in the forward method.

Overall, nn.ModuleList and nn.ModuleDict are useful container classes in Pytorch for managing multiple modules in a neural network model. They provide an organized way to handle complex models with multiple layers or components. You can experiment with these container classes in your own Pytorch projects to see the benefits they offer in terms of code readability and maintainability.