Neural Network Fundamentals
nn.Module, the layers that make up a network, and why a nonlinearity between them is what makes depth mean anything
Neural Network Fundamentals
TL;DR
nn.Module is PyTorch's building block for every layer and every model — you define what parameters exist in __init__ and how data flows through them in forward(). Nonlinear activation functions between layers are not decoration: without one, any number of stacked linear layers collapses mathematically into a single linear layer, and depth stops meaning anything.
| Property | Value |
|---|---|
| Level | Beginner |
| Reading time | ~18 minutes |
| Prerequisites | What Is Deep Learning? |
| You will understand | How to build a network in PyTorch, and why nonlinearity is mathematically necessary |
nn.Module: The Building Block
Every layer and every model in PyTorch is an nn.Module. Two methods matter:
import torch.nn as nn
class TinyClassifier(nn.Module):
def __init__(self, input_dim, hidden_dim, num_classes):
super().__init__()
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.activation = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, num_classes)
def forward(self, x):
return self.fc2(self.activation(self.fc1(x)))
model = TinyClassifier(input_dim=784, hidden_dim=128, num_classes=10)| Method | Role |
|---|---|
__init__ | Declares every layer the model owns — this is where parameters get created |
forward | Defines how an input actually flows through those layers to produce an output |
Layers declared in __init__ but never used in forward() still hold trainable parameters. PyTorch doesn't know a layer is "unused" — it just sees a module with parameters registered on the model, and the optimizer will happily try to train them even though nothing routes data through them. This is a real, if slightly odd, source of wasted parameters and confusing parameter counts in a model someone else wrote.
Common Layer Types
| Layer | Computes | Typical use |
|---|---|---|
nn.Linear(in, out) | y = xW^T + b — a full matrix multiply plus bias | The default workhorse layer, everywhere |
nn.Embedding(vocab_size, dim) | A lookup table — maps an integer index to a learned vector | Turning token IDs into vectors at a model's input |
nn.Conv2d | A sliding-window weighted sum over local regions of an image | Vision models — briefly relevant here for context, covered in depth elsewhere |
embed = nn.Embedding(num_embeddings=50000, embedding_dim=512)
token_ids = torch.tensor([15, 302, 9981])
vectors = embed(token_ids) # shape: (3, 512)An nn.Embedding layer's "weights" are just its lookup table — a (vocab_size, dim) matrix. Looking up a token's embedding is literally indexing into that matrix, which is also why the embedding table can be a disproportionately large fraction of a small model's total parameters, a point that matters a great deal for small language models specifically.
Why Nonlinearity Is Mathematically Necessary
Here's the argument in full, because it's the single most important fact in this page. A linear layer computes y = xW + b. Stack two of them with nothing in between:
y = (xW1 + b1)W2 + b2
= x(W1 W2) + (b1 W2 + b2)
= xW' + b' where W' = W1 W2 and b' = b1 W2 + b2The result is still just a linear function of x — a different W' and b', but structurally identical to a single layer. No number of stacked linear layers can represent anything a single linear layer can't. Depth only adds real representational power once a nonlinear function sits between the layers, breaking that algebraic collapse.
class TwoLayerReal(nn.Module):
def __init__(self, d):
super().__init__()
self.fc1 = nn.Linear(d, d)
self.act = nn.GELU() # <- this is what makes depth matter
self.fc2 = nn.Linear(d, d)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))Activation Functions
ReLU vs GELU
ReLU: max(0, x)
Simple, cheap, and historically the default. Its problem: the gradient is exactly zero for every negative input — a unit that ends up always receiving negative input ("dead ReLU") stops learning entirely, since no gradient signal reaches it.
GELU and similar smooth variants
RecommendedA smooth curve that behaves like ReLU for large positive/negative inputs but transitions gradually near zero instead of having a hard kink. The smoother gradient near zero is part of why GELU (and variants like SwiGLU) became the standard choice in transformer architectures — better gradient flow through very deep stacks of layers.
import torch
x = torch.linspace(-3, 3, 7)
print(torch.relu(x)) # negative values clipped to exactly 0
print(torch.nn.functional.gelu(x)) # negative values smoothly attenuated, not zeroedCounting Parameters, By Hand
A skill worth having directly, not just from model.parameters(): a linear layer's parameter count is (in_features * out_features) + out_features (the + out_features is the bias term).
nn.Linear(784, 128): 784 * 128 + 128 = 100,480 parameters
nn.Linear(128, 10): 128 * 10 + 10 = 1,290 parameters
TinyClassifier total: 100,480 + 1,290 = 101,770 parameterstotal = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)sum(p.numel() for p in model.parameters()) counts every registered parameter, including ones that might be frozen (requires_grad=False) — such as a frozen base model in a LoRA setup. Use the if p.requires_grad filter whenever the number you actually want is "what will the optimizer update," not "how big is this model on disk."
Moving a Model Between Devices
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
inputs = inputs.to(device)A model and its inputs must be on the same device, or PyTorch raises an error rather than silently doing something wrong. .to(device) on a tensor returns a new tensor on that device rather than moving it in place — a common early mistake is calling tensor.to(device) without reassigning the result and then being confused why the tensor is still on the CPU. model.to(device), by contrast, does move the model's parameters in place, which is a real, easy-to-miss asymmetry between the two.
Concept Checks
Check yourself
Why can't stacking ten nn.Linear layers, with no activation function between them, represent anything a single nn.Linear layer can't?
Because composing linear transformations algebraically collapses into a single linear transformation — the product of the weight matrices is just another matrix, and the sum of the bias terms is just another vector. Regardless of how many linear layers are stacked, the overall function remains linear unless a nonlinearity breaks that collapse.
A layer is declared in __init__ but never called in forward(). Does it affect training?
Yes — PyTorch registers its parameters as part of the model regardless of whether forward() actually routes data through it, and the optimizer will update those parameters if they require gradients, even though they have no effect on the model's output. This wastes compute and memory and can produce a misleadingly large parameter count for what the model is actually doing.
Why does sum(p.numel() for p in model.parameters()) sometimes overstate the number of parameters actually being trained?
Because it counts every registered parameter regardless of whether it's frozen — a common case being a frozen base model with a small set of trainable LoRA adapter weights. Filtering with if p.requires_grad gives the count that actually reflects what the optimizer will update, which is usually the number that matters for reasoning about training cost.
Key Concepts Recap
| Concept | One-line summary |
|---|---|
nn.Module | PyTorch's base class — __init__ declares layers, forward defines data flow |
nn.Linear | y = xW^T + b, the default workhorse layer |
nn.Embedding | A learned lookup table from integer index to vector |
| The linear-collapse argument | Stacked linear layers with no nonlinearity between them reduce to one linear layer |
| ReLU vs GELU | GELU's smoother gradient near zero is why it's standard in transformers |
| Parameter counting | in_features * out_features + out_features for a linear layer |
requires_grad filter | Needed to count only trainable, not merely present, parameters |
.to(device) | Moves a model's parameters in place; returns a new tensor for a tensor (must reassign) |
Next
With a network's structure in place, the next page covers how it actually learns: Training Loops & Optimization.
What Is Deep Learning?
A neural network is a parameterized function fit to data by gradient descent — not something that understands. The idea from zero, with a worked example.
Training Loops & Optimization
Backpropagation, the training loop step by step, and how SGD, Adam, and AdamW actually differ — including the bug AdamW specifically fixed