DVP IBM AI Engineering
Reference

Primer 4 — Shapes & the Training Loop

Matrix shapes, broadcasting, and the five-step loop built from scratch

Do this before Course 4. Goal: the five-line PyTorch training loop stops being an incantation, because you've built each piece in plain Python + NumPy first.

Shape errors are the whole game

Most "my code is broken" moments in Courses 4–6 are one of these three. Learn to read them now:

import numpy as np
A = np.array([[1., 2., 3.],
              [4., 5., 6.]])        # shape (2, 3)
B = np.array([[10., 20.],
              [30., 40.],
              [50., 60.]])          # shape (3, 2)

print((A @ B).shape)                # (2, 2) — works: inner numbers match (2,3)@(3,2)
# print((B @ B).shape)              # ERROR — (3,2)@(3,2): 2 ≠ 3

Matrix multiply rule: (m, n) @ (n, p) → (m, p). The two inner numbers must match; the outer two are the answer's shape. When PyTorch says mat1 and mat2 shapes cannot be multiplied (2x3 and 2x3) — it's telling you the inner numbers, 3 and 2, disagree. Not mysterious. Arithmetic.

Broadcasting — the helpful magic that sometimes bites

x = np.array([[1.], [2.], [3.]])    # shape (3, 1) — a column
y = np.array([10., 20., 30.])       # shape (3,)   — a flat vector

print((x + y).shape)                # (3, 3) !!!

NumPy stretched the column against the row and made a 3×3 grid. Sometimes that's exactly what you want. Sometimes — as the Course 4 notes warn about MSELoss with a (N,) target vs (N,1) prediction — it silently computes the wrong thing with no error. That warning is this exact mechanism.

The defence is one habit: print .shape before you trust any operation. (3,) and (3,1) are different animals; .view(-1, 1) / .reshape(-1, 1) turns the flat one into a column ( -1 means "you figure out this dimension").

Self-check — answer without looking
  1. (4, 2) @ (2, 7) — legal? Result shape?
  2. Why is (5,) + (5,1) dangerous rather than illegal?

The training loop, built from scratch

Course 2 taught the theory: forward → loss → backward → update. Here it is with no framework at all — fitting y = w·x to data, in the Code Lab, right now:

import numpy as np

# data that secretly follows y = 3x
x = np.array([1., 2., 3., 4.])
y = np.array([3., 6., 9., 12.])

w = 0.0        # our guess, starting wrong
lr = 0.05      # learning rate: how big each correction step is

for epoch in range(20):
    y_hat = w * x                          # 1. FORWARD: predict with current w
    loss = ((y_hat - y) ** 2).mean()       # 2. LOSS: how wrong, on average (MSE)
    grad = (2 * (y_hat - y) * x).mean()    # 3. BACKWARD: slope of loss wrt w
    w = w - lr * grad                      # 4. STEP: nudge w downhill
    print(f"epoch {epoch:2d}  w={w:.3f}  loss={loss:.3f}")

Run it. Watch w walk from 0.0 to ~3.0 while loss collapses. That's all training is — these four lines, scaled up to millions of parameters.

New syntax decoded:

Now map it onto the real thing

The certificate's five-step PyTorch loop, side by side with what you just built:

for x, y in loader:                # batches, not the whole set (mini-batch GD)
    optimizer.zero_grad()          # clear old gradients — they ACCUMULATE otherwise
    y_hat = model(x)               # 1. FORWARD          (your  w * x)
    loss = criterion(y_hat, y)     # 2. LOSS             (your  MSE line)
    loss.backward()                # 3. BACKWARD         (your  grad line, automated)
    optimizer.step()               # 4. STEP             (your  w = w - lr*grad)

One genuinely new thing: zero_grad(). PyTorch adds new gradients onto old ones (useful for advanced tricks, a trap for everyone else), so you wipe the slate each batch. Your NumPy version recomputed grad fresh each loop — that's why it didn't need the wipe. The Course 4 notes call forgetting this "the classic PyTorch bug"; now you know what it actually means.

Self-check — answer without looking
  1. In your NumPy loop, which line corresponds to loss.backward()?
  2. Set lr = 1.2 and rerun. What happens to loss, and which page of the notes predicted it?
  3. Why does the PyTorch loop need zero_grad() when the NumPy one didn't?

Axis — the last shape idea

mean() averaged everything. Real code averages along a direction:

X = np.array([[1., 2., 3.],
              [4., 5., 6.]])       # (2, 3): 2 samples, 3 features

print(X.mean())          # 3.5        — everything
print(X.mean(axis=0))    # [2.5 3.5 4.5] — down the columns: per-FEATURE mean
print(X.mean(axis=1))    # [2. 5.]       — across the rows:  per-SAMPLE mean

Memory hook: axis=n is the axis that disappears. axis=0 collapses the rows away, leaving one value per column. Feature scaling in Course 1 (StandardScaler) is exactly (X - X.mean(axis=0)) / X.std(axis=0) — you can now read that.

Self-check — answer without looking

scores.shape == (32, 10) — 32 images, 10 class scores each. Which axis do you soft-max over, and why?

Primer 4 in one page