Initialisation, dropout, batch norm, momentum and CNNs — the techniques that make a deep network actually train
21 hours · 6 modules · PyTorch Practice: CNN for Anime Image Classification. Final: Fashion MNIST Classification.
Course 4 stopped at logistic regression. This one goes from there to deep CNNs, and covers the training tricks that make deep networks actually work: initialisation, dropout, batch norm, momentum.
Deliberate overlap with Course 4 Module 6 — same two labs. If you did them, skim.
The argument in one paragraph: with a sigmoid output, MSE produces a non-convex loss surface with plateaus. When the model is confidently wrong, z is large in magnitude, σ'(z) ≈ 0, and the MSE gradient — which contains σ'(z) as a factor — vanishes. The model learns slowest exactly when it is most wrong.
Cross-entropy's derivative cancels that σ'(z) term:
∂J/∂w = (ŷ − y) · x
Gradient proportional to the error. Big mistake, big correction. That's the whole reason cross-entropy is standard for classification.
Run both labs and put the two loss curves side by side — MSE plateaus, cross-entropy descends cleanly.
Generalises sigmoid to K mutually exclusive classes:
softmax(z)ᵢ = e^(zᵢ) / Σⱼ e^(zⱼ)
Outputs are positive and sum to 1 — a probability distribution. Prediction is argmax.
Softmax over 2 classes is algebraically equivalent to sigmoid over 1 output. Different parameterisation, same model.
Numerical stability: e^1000 overflows. Real implementations subtract the max logit first — softmax(z) = softmax(z − max(z)) — which is mathematically identical and safe. PyTorch does this for you only if you use the built-in functions, which is another reason not to hand-roll it.
model = nn.Linear(input_dim, num_classes) # raw logits out
criterion = nn.CrossEntropyLoss() # LogSoftmax + NLLLoss internally
optimizer = optim.SGD(model.parameters(), lr=0.1)
loss = criterion(logits, y) # y = integer class indices, NOT one-hot
_, yhat = torch.max(logits, 1) # predicted class per row
Restating the trap because it's the most-failed item in the course:
Do not put
nn.Softmaxin your model when usingnn.CrossEntropyLoss. The loss applies LogSoftmax itself. Double softmax doesn't error — it just flattens your gradients and the model underperforms for no visible reason.
torch.max(logits, 1) returns (values, indices); you want the indices.
train = datasets.MNIST(root='./data', train=True, download=True,
transform=transforms.ToTensor())
loader = DataLoader(train, batch_size=100, shuffle=True)
model = nn.Linear(28*28, 10)
for x, y in loader:
optimizer.zero_grad()
loss = criterion(model(x.view(-1, 28*28)), y) # flatten images
loss.backward()
optimizer.step()
x.view(-1, 28*28) flattens (batch, 1, 28, 28) → (batch, 784). A linear layer can't take an image tensor directly.
Accuracy ceiling for plain softmax regression on MNIST: roughly 92%. Adding hidden layers (Module 3) gets ~98%; a CNN (Module 5) gets ~99%+. Watching that number climb is the arc of the course.
| Range | Derivative max | Problem | |
|---|---|---|---|
| Sigmoid | (0,1) | 0.25 | vanishing gradient, not zero-centred |
| Tanh | (−1,1) | 1.0 | still saturates at the tails |
| ReLU | [0,∞) | 1 | dying ReLU |
torch.sigmoid(x); torch.tanh(x); torch.relu(x)
nn.Sigmoid(); nn.Tanh(); nn.ReLU(); nn.LeakyReLU(0.01)
The lab compares all three on MNIST — ReLU converges fastest and highest. That empirical result is the point of the exercise.
nn.CrossEntropyLoss do internally, and what must your model NOT do?One hidden layer between input and output. The hidden layer's job is to build a new representation in which the problem becomes linearly separable.
class Net(nn.Module):
def __init__(self, D_in, H, D_out):
super().__init__()
self.linear1 = nn.Linear(D_in, H)
self.linear2 = nn.Linear(H, D_out)
def forward(self, x):
x = torch.sigmoid(self.linear1(x)) # hidden activation
return self.linear2(x) # raw logits out
The classic demonstration: XOR is not linearly separable, so logistic regression fails on it. Add one hidden layer with 2 units and it's solved. Each hidden neuron learns a line; the output layer combines them into a region.
More neurons = more decision boundaries = more capacity = more overfitting risk.
The videos walk 1 → 2 → 6 hidden units on the same data. Too few underfits (can't carve enough regions); too many fits the noise. There's no formula — it's a validated hyperparameter.
model = Net(D_in=2, H=6, D_out=1) # 2 features, binary
model = Net(D_in=784, H=100, D_out=10) # MNIST, 10 classes
Parameter count for a layer: in × out + out. For 784→100: 78,400 + 100 = 78,500. They ask this.
model = nn.Sequential(
nn.Linear(784, 100), nn.ReLU(),
nn.Linear(100, 10),
)
Equivalent to the subclass above for simple stacks. Use the subclass when the forward pass needs logic.
For a two-layer net, the chain rule cascade:
∂J/∂W₂ = δ₂ · a₁ᵀ where δ₂ = (ŷ − y)
∂J/∂W₁ = δ₁ · xᵀ where δ₁ = (W₂ᵀ δ₂) ⊙ f'(z₁)
The f'(z₁) factor is the vanishing gradient in algebraic form: with sigmoid it's ≤ 0.25 and shrinks δ at every layer going backwards.
PyTorch computes all of this from loss.backward(). You need the intuition, not the manual derivation — but the exam does ask which term causes vanishing gradients.
nn.Linear(100, 50)? (5,050)The most practically valuable module in the course. These are the techniques that separate a network that trains from one that doesn't.
class DeepNet(nn.Module):
def __init__(self, layers): # e.g. [784, 100, 50, 10]
super().__init__()
self.hidden = nn.ModuleList()
for in_size, out_size in zip(layers, layers[1:]):
self.hidden.append(nn.Linear(in_size, out_size))
def forward(self, x):
L = len(self.hidden)
for i, layer in enumerate(self.hidden):
x = torch.relu(layer(x)) if i < L - 1 else layer(x)
return x
Use nn.ModuleList, not a plain Python list. A plain list doesn't register the layers as submodules — model.parameters() returns nothing, the optimiser has nothing to update, and training silently does nothing. Silent failure, hard to spot.
During training, randomly zero each neuron with probability p. At inference, keep everything and scale appropriately (PyTorch uses inverted dropout, so it scales during training and does nothing at eval).
self.drop = nn.Dropout(p=0.5)
x = self.drop(torch.relu(self.linear1(x)))
Why it regularises: no neuron can rely on any specific other neuron being present, so the network can't build fragile co-adapted chains. Equivalent to training an ensemble of exponentially many thinned networks and averaging them.
model.eval() disables dropout. Forgetting it means random predictions at inference.This is the single most common cause of "my model works in training and is garbage at test time" in these labs.
Why not zeros: every neuron in a layer computes the same thing, gets the same gradient, and stays identical forever. The layer has the effective capacity of one neuron. This is the symmetry problem, and the "Initialization with Same Weights" lab exists to show it.
Why not large random: activations saturate, gradients vanish. Why not small random: activations shrink toward zero through depth, signal dies.
The goal is to keep activation variance roughly constant across layers.
| Method | Variance | Use with |
|---|---|---|
| Xavier / Glorot | 2/(n_in + n_out) | sigmoid, tanh |
| He / Kaiming | 2/n_in | ReLU |
| Uniform | fixed range | baseline for comparison |
He accounts for ReLU zeroing half the activations, so it uses twice the variance Xavier would.
nn.init.xavier_uniform_(layer.weight)
nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu')
The labs compare default / Xavier / He on MNIST with tanh and with ReLU. The pairing is the lesson: Xavier with tanh, He with ReLU.
v := ρ·v + ∂J/∂w
w := w − α·v
Accumulate a velocity vector. Consistent gradient directions build speed; oscillating directions cancel out. ρ (momentum) is typically 0.9.
Three things it fixes:
the zigzag and accelerates along the floor
optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
The "Momentum with Different Polynomials" lab visualises this on a 1-D loss surface with a local minimum — with ρ=0 you get stuck, with ρ=0.9 you roll through.
Normalise each layer's pre-activations across the batch, then learn a scale and shift:
x̂ = (x − μ_batch) / √(σ²_batch + ε)
y = γ·x̂ + β γ, β are LEARNED
The learnable γ and β matter: normalising to mean 0 / variance 1 would remove the layer's ability to use the non-linear parts of its activation function. γ and β let the network undo the normalisation if that's what's best.
self.bn1 = nn.BatchNorm1d(hidden_size) # 1d for dense layers
x = torch.relu(self.bn1(self.linear1(x)))
nn.BatchNorm2d(num_channels) # 2d for conv layers
Benefits: faster convergence, tolerates higher learning rates, reduces sensitivity to initialisation, adds slight regularisation (batch statistics are noisy).
Train vs eval: BatchNorm uses batch statistics during training and running averages at inference. model.eval() switches the mode. This is the second reason train()/eval() are mandatory — and with a batch size of 1 BatchNorm can't compute a variance at all.
Conventional order: Linear → BatchNorm → Activation → Dropout. (BatchNorm-then-activation is what the original paper and these labs use. Some later work argues the reverse. Follow the course.)
model.eval() changes, and the symptom of forgetting it.Slide a small learned kernel over the input, computing a dot product at each position. Output is an activation map / feature map.
Three properties, all examinable:
regardless of image size
Output size formula — you will be asked to compute this:
out = floor( (in − kernel + 2·padding) / stride ) + 1
Worked: input 28, kernel 5, padding 0, stride 1 → (28−5)/1 + 1 = 24. With padding 2: (28−5+4)/1 + 1 = 28 (unchanged — that's "same" padding).
under-sampled. padding = (kernel−1)/2 keeps size constant for odd kernels.
nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2)
Downsample. Max pooling takes the maximum in each window.
nn.MaxPool2d(kernel_size=2, stride=2) # halves height and width
Gives small-translation invariance, cuts compute, enlarges the effective receptive field. No learnable parameters. Max pooling is standard; average pooling is used mainly as global pooling before the classifier head.
An RGB image has 3 input channels. A conv layer with out_channels=16 learns 16 separate kernels, each spanning all input channels, producing 16 feature maps.
Parameters in a conv layer:
(kernel_h × kernel_w × in_channels × out_channels) + out_channels
For Conv2d(3, 16, 5): 5×5×3×16 + 16 = 1,216. A dense layer on a 224×224×3 image with 16 outputs would need 2.4 million. That gap is why CNNs exist.
class CNN(nn.Module):
def __init__(self, out_1=16, out_2=32, num_classes=10):
super().__init__()
self.cnn1 = nn.Conv2d(1, out_1, kernel_size=5, padding=2)
self.bn1 = nn.BatchNorm2d(out_1)
self.pool = nn.MaxPool2d(2)
self.cnn2 = nn.Conv2d(out_1, out_2, kernel_size=5, padding=2)
self.bn2 = nn.BatchNorm2d(out_2)
self.fc = nn.Linear(out_2 * 7 * 7, num_classes)
def forward(self, x):
x = self.pool(torch.relu(self.bn1(self.cnn1(x)))) # 28→28→14
x = self.pool(torch.relu(self.bn2(self.cnn2(x)))) # 14→14→7
x = x.view(x.size(0), -1) # flatten
return self.fc(x)
The out_2 * 7 * 7 is where everyone gets stuck: you must trace the spatial dimensions through every layer to know the flattened size. 28 →(pool) 14 →(pool) 7. With 32 channels that's 32×7×7 = 1568.
x.view(x.size(0), -1) — keep the batch dimension, flatten the rest.
Deeper networks were observed to perform worse than shallower ones — not from overfitting but because the optimiser couldn't train them. ResNet's fix: skip connections.
output = F(x) + x
The layer learns a residual F(x) = output − x rather than the full mapping. If the identity is optimal, F just has to learn zero, which is easy. And the + x gives the gradient a direct path backwards, so it can't vanish through depth.
This one idea is why 50-, 101-, 152-layer networks became trainable — and it's the same residual pattern inside every transformer block.
from torchvision import models
model = models.resnet18(pretrained=True)
for param in model.parameters():
param.requires_grad = False # freeze everything
model.fc = nn.Linear(512, num_classes) # replace the head (new layers default to
# requires_grad=True)
optimizer = optim.Adam(model.fc.parameters(), lr=1e-3) # optimise ONLY the head
Two details the lab checks: replace model.fc (ResNet) or model.classifier (VGG/DenseNet) depending on architecture, and pass only the trainable parameters to the optimiser.
Pretrained torchvision models expect: 3-channel RGB, at least 224×224, normalised with mean [0.485, 0.456, 0.406] and std [0.229, 0.224, 0.225]. MNIST is 1-channel 28×28 — you must resize and repeat the channel.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
for x, y in loader:
x, y = x.to(device), y.to(device) # BOTH model and data must be on the device
RuntimeError: Expected all tensors to be on the same device means you moved one but not the other. Move tensors back with .cpu() before converting to NumPy.
Conv2d(3, 8, kernel_size=3, padding=1), then MaxPool2d(2). Output shape?(8×16×16)
Conv2d(16, 32, kernel_size=3)? (4,640)Practice: CNN for Anime Image Classification. Final: Fashion MNIST Classification — 10 clothing classes, 28×28 grayscale.
Rubric:
nn.Module with conv → pool → FC, architecture justifiedCrossEntropyLoss) and optimiserReference implementation:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.2860,), (0.3530,)), # FashionMNIST stats
])
train = datasets.FashionMNIST('./data', train=True, download=True, transform=transform)
test = datasets.FashionMNIST('./data', train=False, download=True, transform=transform)
model = CNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(n_epochs):
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
model.eval()
correct = 0
with torch.no_grad():
for x, y in test_loader:
x, y = x.to(device), y.to(device)
correct += (model(x).argmax(1) == y).sum().item()
print(f'epoch {epoch}: {correct/len(test):.4f}')
Fashion-MNIST reference accuracies: plain softmax ~84%, MLP ~88%, small CNN ~91–93%, CNN with batchnorm + augmentation ~93–94%. Shirt/coat/pullover is where the confusion matrix always concentrates — say so in your write-up, it's exactly the kind of analysis that scores.
He init→ReLU signal, momentum→ravines and local minima, residuals→depth
train() vs eval() effectsCrossEntropyLossnn.ModuleList for dynamic layer lists — a plain list won't register parameterseval()out = (in − k + 2p)/s + 1 — know it coldF(x) + x gives gradients a highway; makes very deep nets trainable