Common layers: Linear, Conv2d, activations and Sequential

The layers you will use most

nn.Linear(in_features, out_features) computes y = x W^T + b on the last dimension, so it accepts (N, in_features) or (N, T, in_features). nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding) slides learned filters over images shaped (N, C, H, W); its output height is floor((H + 2*padding - kernel) / stride) + 1. nn.MaxPool2d(2) halves height and width. nn.Flatten() turns (N, C, H, W) into (N, C*H*W) before a Linear layer.

Activations such as nn.ReLU and nn.GELU add non-linearity. Without them, any stack of Linear layers collapses into a single linear function, however deep it is.

nn.Dropout(p) randomly zeroes activations during training to reduce overfitting. nn.BatchNorm2d and nn.LayerNorm normalise activations, which usually makes training faster and more stable. nn.Sequential(...) chains modules in order and is perfect for simple stacks; write a custom forward when you need branches or skip connections.

A tiny CNN with nn.Sequential

The comments track the shape after each layer. Getting 16 * 14 * 14 right is the usual sticking point; computing it from the conv formula avoids guesswork.

import torch
from torch import nn

model = nn.Sequential(
    nn.Conv2d(1, 16, kernel_size=3, padding=1),  # (N,1,28,28) -> (N,16,28,28)
    nn.ReLU(),
    nn.MaxPool2d(2),                             # -> (N,16,14,14)
    nn.Flatten(),                                # -> (N, 3136)
    nn.Linear(16 * 14 * 14, 10),                 # -> (N, 10)
)
x = torch.randn(8, 1, 28, 28)
print(model(x).shape)    # torch.Size([8, 10])

Convolution output size in plain Python

Run this whenever you design a conv stack. Padding 1 with a 3x3 kernel keeps the size; no padding shrinks it; stride 2 roughly halves it.

def conv_out(size, kernel, stride=1, padding=0):
    return (size + 2 * padding - kernel) // stride + 1

print(conv_out(28, 3, padding=1))   # padding 1 keeps 28
print(conv_out(28, 5))              # no padding shrinks
print(conv_out(224, 7, stride=2, padding=3))

Output:

28
24
112

Photocopiers without creativity

Stacking Linear layers without activations is like passing a document through five photocopiers that each only resize and shift it. The final copy is still just one resize and shift of the original; one machine could have done it. An activation is like a copier that can fold the paper, and folds can combine into shapes no single resize can make.

Quick check: Why are non-linear activations placed between Linear layers?

• They reduce the number of parameters
• Without them, stacked Linear layers are equivalent to a single Linear layer
• They are required for autograd to compute gradients
• They convert logits into probabilities