The ‘nn.Module‘ and ‘nn.Sequential‘ classes are fundamental building blocks in PyTorch, but they serve different purposes.
‘nn.Module‘ is the base class for all neural network modules in PyTorch. It has many useful features, such as the ability to track all its parameters, to register sub-modules and to implement various hooks to manipulate the forward and backward computation of modules. Using ‘nn.Module‘ is essential when designing complex deep learning models that require multiple layers, branches, or concatenation.
Here’s an example of how to use ‘nn.Module‘ to define a two-layer neural network:
import torch.nn.functional as F
import torch.nn as nn
class TwoLayerNet(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(TwoLayerNet, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.fc2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
In this example, we define a two-layer neural network with ‘nn.Linear‘ and ‘F.relu‘ activation function. All the layers are registered to the module as attributes (‘self.fc1‘ and ‘self.fc2‘) inside the ‘__init__‘ method. Then, the forward pass is implemented inside the ‘forward‘ method, which applies the layers in the desired order to the input data.
‘nn.Sequential‘, on the other hand, is a container module that allows you to stack multiple layers in a sequence. It is a simple way to construct a neural network, especially if the network consists of a linear sequence of layers. It can be considered as a shortcut to ‘nn.Module‘ that saves some boilerplate code when defining neural networks.
Here’s an example of how to use ‘nn.Sequential‘ to define the same two-layer neural network:
model = nn.Sequential(
nn.Linear(input_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, num_classes)
)
In this example, we define the same two-layer neural network, but we use ‘nn.Sequential‘ to stack the layers one after the other with a specified activation function between them.
In conclusion, both ‘nn.Module‘ and ‘nn.Sequential‘ are important building blocks in PyTorch. While ‘nn.Module‘ is essential for designing complex and intricate models with multiple parts, ‘nn.Sequential‘ is a convenient way of stacking layers in a linear fashion.