Why nn.Module organizes model code in PyTorch

Introduction

nn.Module helps keep your model code neat and easy to manage. It groups all parts of your model together so you can build, train, and save it smoothly.

When building a neural network with PyTorch.

When you want to reuse parts of your model easily.

When you need to save or load your model weights.

When you want to keep your training code clean and organized.

Syntax

PyTorch

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # define layers here

    def forward(self, x):
        # define forward pass here
        return x

Always call super().__init__() to initialize the base class.

Define your layers in __init__ and the data flow in forward.

Examples

A simple model with one linear layer.

PyTorch

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

A model with two layers and a ReLU activation in between.

PyTorch

class TwoLayerNN(nn.Module):
    def __init__(self):
        super(TwoLayerNN, self).__init__()
        self.layer1 = nn.Linear(10, 20)
        self.layer2 = nn.Linear(20, 5)

    def forward(self, x):
        x = self.layer1(x)
        x = nn.functional.relu(x)
        x = self.layer2(x)
        return x

Sample Model

This code defines a simple model using nn.Module, then runs a sample input through it to get the output.

PyTorch

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(3, 1)

    def forward(self, x):
        return self.linear(x)

model = SimpleModel()

# Create a sample input tensor
input_tensor = torch.tensor([[1.0, 2.0, 3.0]])

# Get model output
output = model(input_tensor)

print(f"Input: {input_tensor}")
print(f"Output: {output}")

Important Notes

nn.Module tracks all layers automatically, so you don't have to manage them yourself.

It also helps with moving your model to GPU or CPU easily.

Using nn.Module makes saving and loading models straightforward.

Summary

nn.Module organizes your model code by grouping layers and logic.

It makes training, saving, and reusing models easier.

Always define layers in __init__ and data flow in forward.