0Pricing
Machine Learning Academy · Lesson

Building a Feedforward Network with nn.Module

Learners will subclass nn.Module, stack Linear and ReLU layers in __init__, implement the forward pass, and verify output shapes with a dummy input.

What Is nn.Module?

nn.Module is PyTorch's base class for all neural network components. Every layer, activation function, loss function, and complete model in PyTorch subclasses nn.Module. It provides parameter management, device movement, serialisation, and training/eval mode switching out of the box. By subclassing it you get all this functionality for free and only need to define the architecture and the forward pass.

import torch.nn as nn

# Inspect what nn.Module gives you
model = nn.Linear(4, 2)
print(type(model))               # <class 'torch.nn.modules.linear.Linear'>
print(isinstance(model, nn.Module))  # True
print(list(model.parameters()))  # weight and bias tensors

Subclassing nn.Module: __init__ and forward

Creating a custom network requires two methods. In __init__ you call super().__init__() and define all learnable layers as attributes. In forward you describe how data flows through those layers. PyTorch tracks any nn.Module or nn.Parameter assigned to self as a trainable component. Calling the model like a function (model(x)) automatically invokes forward and runs any registered hooks.

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

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

model = SimpleNet(4, 16, 3)
print(model)

All lessons in this course

  1. PyTorch Tensors: Creation, Operations, and GPU Transfer
  2. Autograd: Automatic Differentiation for Backpropagation
  3. Building a Feedforward Network with nn.Module
  4. Training Loop: Loss, Optimizer, and Epochs
← Back to Machine Learning Academy