Advanced Module Architecture

Beyond Sequential Models

You're already familiar with nn.Sequential for creating straightforward, linear stacks of layers. It's clean and effective for many tasks. But modern neural networks often look less like a straight line and more like a branching tree, with complex data flows and custom logic.

To build these sophisticated architectures, you need to move beyond nn.Sequential and use its parent, nn.Module. By subclassing nn.Module, you gain complete control over your model's structure and behavior. It's the standard way to define professional, reusable, and complex models in PyTorch.

To define a neural network in PyTorch, we create a class that inherits from nn.Module.

Let's see how this works. Every custom model you build will have two essential methods:

__init__(): This is where you define and initialize your model's layers, such as convolutions, linear layers, or even other custom nn.Module instances.
forward(): This method defines the data flow. It takes the input tensor(s) and passes them through the layers you defined in __init__(), returning the final output.

Here’s a basic skeleton for a custom model. We'll define a simple feed-forward network that we could have built with nn.Sequential, but now it's a custom class.

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        # Call the constructor of the parent class (nn.Module)
        super().__init__()
        
        # Define the layers. PyTorch will automatically track their parameters.
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.activation = nn.ReLU()
        self.layer2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # Define the forward pass
        x = self.layer1(x)
        x = self.activation(x)
        x = self.layer2(x)
        return x

# Instantiate the model
model = SimpleNet(input_size=784, hidden_size=256, output_size=10)
print(model)

This structure gives you a blueprint. You define the parts in __init__ and then arrange how data flows through them in forward. This separation of components from their execution is what makes nn.Module so powerful.

Dynamic and Flexible Architectures

What if you don't know the exact number of layers your model needs at initialization? Or what if you want to create a more dynamic structure? Using standard Python lists or dictionaries to hold your layers won't work, because PyTorch won't be able to find and train their parameters.

PyTorch provides two special containers for this: nn.ModuleList and nn.ModuleDict.

nn.ModuleList holds modules in a list-like structure, while nn.ModuleDict uses a dictionary-like format. Both ensure that all sub-modules are properly registered with your parent nn.Module.

Let's build a model with a variable number of hidden layers using nn.ModuleList. This is useful when you want to experiment with network depth without rewriting the class.

class DynamicNet(nn.Module):
    def __init__(self, input_size, output_size, hidden_sizes):
        super().__init__()
        self.layers = nn.ModuleList()
        
        # Create layers dynamically based on the hidden_sizes list
        in_size = input_size
        for h_size in hidden_sizes:
            self.layers.append(nn.Linear(in_size, h_size))
            self.layers.append(nn.ReLU())
            in_size = h_size # The output of this layer is the input to the next
            
        # Add the final output layer
        self.layers.append(nn.Linear(in_size, output_size))

    def forward(self, x):
        # Iterate through the ModuleList to apply each layer
        for layer in self.layers:
            x = layer(x)
        return x

# Create a model with three hidden layers of sizes 128, 64, and 32
model = DynamicNet(784, 10, hidden_sizes=[128, 64, 32])
print(model)

nn.ModuleDict is useful when you want to associate layers with specific names, making your forward pass more readable. This is common in models with distinct, named branches.

class BranchingNet(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.branches = nn.ModuleDict({
            'branch1': nn.Linear(20, 10),
            'branch2': nn.Linear(20, 10)
        })
        
        self.final = nn.Linear(10, 1)

    def forward(self, x, branch_name):
        # Select a branch to use based on the branch_name key
        x = self.branches[branch_name](x)
        return self.final(x)

model = BranchingNet()
input_tensor = torch.randn(1, 20)

# Pass data through a specific branch
output = model(input_tensor, branch_name='branch1')
print(output)

Customizing the Forward Pass

The true power of subclassing nn.Module shines in the forward method. Since it's just a Python function, you can implement any logic you need: loops, conditional statements, and calls to other functions. This allows for architectures that are impossible with nn.Sequential.

A great example is stochastic depth, a regularization technique used in networks with residual connections (like ResNets). During training, it randomly "drops" entire residual blocks, forcing the network to learn more robust features. During evaluation, all blocks are used.

This requires conditional logic: we need to check if the model is in training or evaluation mode.

# A residual block that can be randomly dropped
class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.main_path = nn.Sequential(
            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(channels, channels, kernel_size=3, padding=1)
        )
    
    def forward(self, x):
        return x + self.main_path(x)


# Model implementing stochastic depth
class ResNetWithStochasticDepth(nn.Module):
    def __init__(self, channels, drop_prob=0.5):
        super().__init__()
        self.block = ResidualBlock(channels)
        self.drop_prob = drop_prob

    def forward(self, x):
        # Only apply stochastic depth during training
        if self.training and torch.rand(1).item() < self.drop_prob:
            # Skip the block entirely
            return x
        else:
            # Use the block
            return self.block(x)

# When model.train() is called, stochastic depth is active.
# When model.eval() is called, it's disabled.

The self.training attribute is a boolean flag inherited from nn.Module. It's automatically set to True when you call model.train() and False with model.eval(). This is the standard way to implement behaviors that should only happen during training, like dropout or stochastic depth.

Initialization and State

How you initialize your model's weights can have a huge impact on training stability and convergence speed. Poor initialization can lead to vanishing or exploding gradients. While PyTorch provides reasonable default initializations, for advanced models, custom strategies are often necessary.

A common practice is to create a dedicated method within your nn.Module to handle weight initialization and apply it recursively to all sub-modules.

class ModelWithCustomInit(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(100, 50)
        self.layer2 = nn.Linear(50, 10)
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                # Kaiming He initialization for layers followed by ReLU
                nn.init.kaiming_uniform_(m.weight, nonlinearity='relu')
                if m.bias is not None:
                    # Initialize bias to zero
                    nn.init.constant_(m.bias, 0)
            
    def forward(self, x):
        return self.layer2(self.layer1(x))

The self.modules() method returns an iterator over all modules in the network, allowing you to apply initialization logic to each one. Here, we used Kaiming initialization, which is well-suited for layers followed by a ReLU activation.

Finally, let's distinguish between two ways nn.Module can store state: parameters and buffers.

Type	Description	Gradient Tracking	Part of `state_dict`
Parameter	A tensor that is considered a model parameter to be trained.	Yes (requires_grad=True)	Yes
Buffer	A stateful tensor that is part of the model but not a parameter.	No (requires_grad=False)	Yes

You've already been using parameters. Any nn.Module (like nn.Linear) assigned as an attribute is automatically registered. To register a raw tensor as a parameter, you wrap it in nn.Parameter.

Buffers are for state that you want to save with the model's state_dict, but that shouldn't be trained. A common example is the running mean and variance in a BatchNorm layer. You can register your own buffer with register_buffer().

class ModelWithBuffer(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 10)
        
        # Register a buffer to count how many times forward is called
        self.register_buffer('forward_count', torch.tensor(0, dtype=torch.long))
        
        # Register a parameter manually
        self.my_param = nn.Parameter(torch.randn(5))

    def forward(self, x):
        # Buffers can be modified in-place
        self.forward_count += 1
        return self.layer(x)

model = ModelWithBuffer()
print(model.state_dict())

Notice how both forward_count (the buffer) and my_param (the parameter) are included in the state_dict. However, only my_param and the weights of self.layer will have gradients computed during backpropagation. This distinction is crucial for managing complex model state effectively.

Quiz Questions 1/5

When creating a custom neural network by subclassing nn.Module, which two methods are essential to implement?

Quiz Questions 2/5

Why must you use nn.ModuleList or nn.ModuleDict instead of a standard Python list or dictionary to hold a dynamic number of layers in a custom nn.Module?

Advanced PyTorch Model Development