Building Neural Networks using nn.Module in PyTorch

Introduction

nn.Module is the base class for all neural networks in PyTorch. It helps you organize layers, parameters, and the forward pass in a clean, modular way. In this page, we’ll define a simple feedforward neural network using nn.Module.

Defining a Custom Neural Network

The custom network class must inherit from nn.Module and implement:

• __init__() – for defining layers.
• forward() – for implementing forward propagation logic.

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(4, 3)  # Input layer
        self.relu = nn.ReLU()      # Activation function
        self.fc2 = nn.Linear(3, 1)  # Output layer

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

Creating and Using the Model

model = SimpleNet()
print(model)

# Input sample
input_data = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
output = model(input_data)
print("Output:", output)

Accessing Model Parameters

All weights and biases of layers can be accessed and updated using:

for name, param in model.named_parameters():
    print(name, param.shape)

Loss Function and Optimizer

To train the model, you need a loss function and an optimizer:

criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

Training Loop Example

# Dummy target
target = torch.tensor([[1.0]])

# Forward pass
output = model(input_data)
loss = criterion(output, target)

# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()

print("Loss:", loss.item())

Extended: Multi-Layer Feedforward Network

import torch
import torch.nn as nn

class MultiLayerNet(nn.Module):
    def __init__(self):
        super(MultiLayerNet, self).__init__()
        self.fc1 = nn.Linear(4, 6)
        self.act1 = nn.Tanh()
        self.fc2 = nn.Linear(6, 3)
        self.act2 = nn.ReLU()
        self.fc3 = nn.Linear(3, 1)

    def forward(self, x):
        x = self.act1(self.fc1(x))
        x = self.act2(self.fc2(x))
        x = self.fc3(x)
        return x

model = MultiLayerNet()
input_data = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
output = model(input_data)

print(output)

• Includes two hidden layers and two different activation functions (Tanh and ReLU).
• More complex structure than the basic example, useful for capturing nonlinear relationships.
• Output is a scalar prediction from the final linear layer.

Using Batch Inputs for Inference

batch_input = torch.tensor([
    [1.0, 2.0, 3.0, 4.0],
    [4.0, 3.0, 2.0, 1.0]
])
batch_output = model(batch_input)

print(batch_output)

• This demonstrates model inference on multiple inputs simultaneously (batch size = 2).
• PyTorch models automatically process batches if the input has shape (batch_size, features).

Customizing Parameter Initialization

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        nn.init.zeros_(m.bias)

model.apply(init_weights)

• Apply Xavier uniform initialization to weights for better convergence in deep models.
• apply() runs the function recursively on all submodules (like Linear layers).

Use Case: Regression with nn.Module

# Predicting a numeric output (e.g., house price)
import torch.optim as optim

model = MultiLayerNet()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

X = torch.tensor([
    [1200.0, 3.0, 2.0, 1.0]
])
y = torch.tensor([[250000.0]])

# Training step
output = model(X)
loss = criterion(output, y)

optimizer.zero_grad()
loss.backward()
optimizer.step()

print("Loss:", loss.item())

• Model predicts a single value like price based on input features (area, rooms, etc.).
• Uses Adam optimizer and MSE loss for regression training.

Freezing Layers in nn.Module

for param in model.fc1.parameters():
    param.requires_grad = False

• Freezing is helpful when using pre-trained layers and only updating specific layers.
• Only the unfrozen layers will contribute to loss.backward().

Checking Trainable Parameters

for name, param in model.named_parameters():
    if param.requires_grad:
        print(name, param.shape)

• Displays only trainable (unfrozen) parameters for model diagnostics or debugging.

Conclusion

Using nn.Module gives structure and clarity to PyTorch model definitions. It allows you to encapsulate model architecture, reuse components, and integrate easily with optimizers and training loops.

Subhendu Mohapatra

Author & Instructor at plus2net

I write and maintain practical tutorials on Python, PHP, SQL, JavaScript, HTML, jQuery, and web development at plus2net. The tutorials focus on clear explanations, working examples, and code that readers can test and adapt while learning.

About the author   Watch on YouTube

Connect: