Neural Network Modules - PyTorch Basics

Neural Network Modules

What you'll build: the three building blocks every PyTorch model is made of (a Linear layer, a ReLU activation and a Sequential container) and a create_mlp factory that chains them into a 2-layer MLP.

Background

In PyTorch every layer is an nn.Module: it owns its parameters, defines forward(x), and can contain other modules. Calling model(x) runs forward, and model.parameters() walks the whole tree to find every trainable tensor, which is what the optimizer receives. Writing these classes yourself shows there is no magic: a model is a tree of objects with weights and a forward method.

Specification

• Linear(in_features, out_features) computes Y = X @ W + b.
  • self.W: shape (in_features, out_features), initialised to randn * 0.01 (small random values break the symmetry between units).
  • self.b: shape (out_features,), initialised to zeros.
  • forward(x): x of shape (batch, in_features) returns (batch, out_features).
• ReLU: forward(x) returns max(0, x) element-wise, same shape.
• Sequential(*layers): forward(x) passes x through the layers in order and returns the last output.
• create_mlp(input_dim, hidden_dim, output_dim) returns Sequential(Linear(input_dim, hidden_dim), ReLU(), Linear(hidden_dim, output_dim)).

A 2-layer MLP has input_dim*hidden_dim + hidden_dim + hidden_dim*output_dim + output_dim parameters: 325 for (10, 20, 5).

Track Differences

	PyTorch track	NumPy track
Base class	subclass nn.Module, call super().__init__()	plain Python classes
Parameters	wrap W and b in nn.Parameter; parameters() is inherited	Linear.parameters() returns [self.W, self.b]
Calling	layer(x) (runs forward)	layer.forward(x)

Constraints

Write the layers yourself: don't wrap nn.Linear, nn.ReLU or nn.Sequential.

Follow-up

• Why must W not be initialised to all zeros?
• Why does PyTorch store nn.Linear.weight as (out_features, in_features) and compute x @ W.T?

Input: create_mlp(10, 20, 5).forward(np.random.randn(4, 10)).shape
Output: (4, 5)

Input: ReLU().forward(np.array([-2.0, 0.0, 1.5]))
Output: array([0. , 0. , 1.5])

import numpy as np

class Linear:
    """Linear layer: Y = X @ W + b"""

    def __init__(self, in_features: int, out_features: int):
        # Create self.W (in_features, out_features) with np.random.randn(...) * 0.01
        # and self.b (out_features,) with zeros
        # Your code here
        pass

    def forward(self, x: np.ndarray) -> np.ndarray:
        """x: (batch, in_features) -> (batch, out_features)"""
        # Your code here
        pass

    def parameters(self):
        """Return the list [W, b]."""
        # Your code here
        pass


class ReLU:
    """ReLU activation: max(0, x)"""

    def forward(self, x: np.ndarray) -> np.ndarray:
        # Your code here
        pass


class Sequential:
    """Container that applies its layers in order"""

    def __init__(self, *layers):
        self.layers = layers

    def forward(self, x: np.ndarray) -> np.ndarray:
        # Your code here
        pass


def create_mlp(input_dim: int, hidden_dim: int, output_dim: int):
    """Create a 2-layer MLP: Linear -> ReLU -> Linear"""
    # Your code here
    pass