Day 05 · Overview
Week 1FNX 101 · 05Superpowers: PyTorch + GPU
Trade your NumPy net for PyTorch and a GPU: autograd, MNIST, and a digit reader that deciphers real handwriting.
Today you will
3- Meet tensors and let autograd compute gradients for you
- Train a real network on MNIST and read thousands of handwritten digits
- Watch a GPU demolish a CPU at the one operation that matters: matrix multiply
Today's items
2 h 30 minLab
FNX 101 · 05Day 5 — Superpowers: PyTorch + GPU ⚡
Mission briefing: Yesterday you hand-built a neural network's engine in pure NumPy — and felt the pain: roughly 130 forward passes just to nudge the weights once. Today the Lab issues you power tools. A library called PyTorch computes those nudges automatically and exactly, and a GPU runs them at ludicrous speed. Your mission: teach a network to read handwriting — then point it at your own.
Previously in the Lab: on Day 4 you built a neural net by hand and earned a deep respect for backpropagation (and no module — just calluses). Today you get the tools the pros actually use.
Today you will:
- Meet tensors — NumPy arrays that can live on a GPU and remember their own history
- Let autograd compute gradients for you (no more finite-difference nudging)
- Train a real network on MNIST and read thousands of strangers' handwritten digits
- Watch a GPU demolish a CPU at the one operation that matters: matrix multiply
- Upgrade to a convolutional network, then draw your own digit for it to read
Today you unlock: 🔓 Digit Reader — draw a digit on screen and the network you trained names it.
A precise reading
MNIST has 60,000 training and 10,000 test images; the test images do not train the model. Autograd applies the chain rule to recorded tensor operations, subject to floating-point precision. GPU speed depends on matrix size, hardware and transfer overhead; small workloads can be faster on the CPU.
# ╔═══════════════════════════════════════════════════════════════════╗
# ║ SETUP — run me first! (identical in every Lab notebook) ║
# ╚═══════════════════════════════════════════════════════════════════╝
import os, sys, pathlib
IN_COLAB = "google.colab" in sys.modules
SMOKE = os.environ.get("SMOKE_TEST", "0") == "1" # tiny fast run for automated tests
REPO_URL = "https://github.com/A-Halimi/ai-studio-internship.git" # instructor: set once
if IN_COLAB:
if not pathlib.Path("/content/ai-studio-internship").exists():
!git clone -q {REPO_URL} /content/ai-studio-internship
%cd /content/ai-studio-internship/notebooks/day05 # ← ADAPT day folder
%pip -q install gradio==6.19.0 # ← ADAPT per-day installs (see spec; delete line if none)
HERE = pathlib.Path.cwd()
REPO = HERE.parents[1] # .../notebooks/dayNN -> repo root
DATA_DIR = pathlib.Path(os.environ.get("COURSE_DATA_DIR", REPO / "data"))
SAMPLES_DIR = REPO / "data_samples" # small datasets shipped with the repo
MODELS_DIR = REPO / "ai_studio" / "models" # where your Studio modules unlock
for p in (REPO / "data", MODELS_DIR):
p.mkdir(parents=True, exist_ok=True)
import random
import numpy as np
SEED = 42
random.seed(SEED); np.random.seed(SEED)
print(f"✅ Setup done | Colab: {IN_COLAB} | Smoke test: {SMOKE}")
print(f" data: {DATA_DIR}")import torch
torch.manual_seed(SEED)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if DEVICE.type == "cuda":
print(f"⚡ Running on GPU: {torch.cuda.get_device_name(0)}")
else:
print("🐢 No GPU found — running on CPU. Everything still works, just slower.")# ── CONFIG — the Lab's control panel. Tweak me later! ──────────────
EPOCHS = 1 if SMOKE else 3
TRAIN_SUBSET = 2_000 if SMOKE else 60_000 # how many training images to use
BATCH_SIZE = 128
LEARNING_RATE = 1e-3
print(f"EPOCHS={EPOCHS} TRAIN_SUBSET={TRAIN_SUBSET:,} BATCH_SIZE={BATCH_SIZE} LR={LEARNING_RATE}")1. Tensors — NumPy arrays with two superpowers
A tensor is PyTorch's version of a NumPy array: an n-dimensional grid of numbers. You already know how to work with these from Days 2–4. Tensors add two things a plain NumPy array can't do:
- They can live on a GPU — thousands of numbers multiplied in parallel.
- They remember their own history — so PyTorch can compute gradients automatically (that's the next section).
The bridge between the two worlds is one function each way: torch.from_numpy(array) and tensor.numpy().
import numpy as np
import torch
a = np.array([[1., 2.], [3., 4.]])
t = torch.from_numpy(a) # numpy -> torch
back = t.numpy() # torch -> numpy
print("numpy array:\n", a)
print("torch tensor:\n", t)
print("shape:", t.shape, "| dtype:", t.dtype)🧪 Exercise 1 — your first tensors
Ais 3×4 andBis 4×2. Matrix-multiply them with the@operator intoC.- Print
C's shape and dtype. - Expected:
Cis 3×2 — the shared inner dimension (4) cancels, exactly like the matmuls you traced by hand on Day 4.
A = torch.rand(3, 4)
B = torch.rand(4, 2)
# ==================== YOUR CODE HERE ====================
### HINT: '@' is matrix multiply; the inner dims (4) must match and cancel
...
print("C shape:", C.shape, "| dtype:", C.dtype)
assert C.shape == (3, 2), "C should be 3x2 — check which dimensions cancel"
print("✅ nice — that's a matrix multiply, the heartbeat of every network")2. Autograd — yesterday's nudge, automatic and exact
On Day 4 you found gradients by nudging: bump a weight a tiny bit, re-run the network, see how the loss changed. It worked, but it was slow and only approximate.
PyTorch does it exactly. Mark a tensor with requires_grad=True and it silently records every operation you do to it. Call .backward() and PyTorch replays those operations in reverse — the chain rule, done for you — leaving the gradient in .grad.
🔢 Math Corner — what backward() actually did (optional)
y = x**2 builds a tiny computation graph: x → square → y. y.backward() walks that graph in reverse, applying the chain rule at each step, and drops the result into x.grad. For y = x² the derivative is 2x, so at x = 2 you get 4. The magic: the same reverse walk works for a graph with a million nodes — which is exactly a neural network.
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
y.backward() # fill in dy/dx
print("y =", y.item())
print("dy/dx at x=2 =", x.grad.item(), " (calculus says 2x = 4 ✓)")🧪 Exercise 2 — differentiate by yourself (well, by PyTorch)
- Build
y = x**3 + 2*xatx = 3.0, then cally.backward(). - Expected:
x.gradis 29. (By hand: d/dx(x³ + 2x) = 3x² + 2 = 3·9 + 2 = 29.)
x = torch.tensor(3.0, requires_grad=True)
# ==================== YOUR CODE HERE ====================
### HINT: write the formula for y, then call y.backward()
...
print("dy/dx at x=3:", x.grad.item())
assert abs(x.grad.item() - 29.0) < 1e-5, "d/dx (x^3 + 2x) = 3x^2 + 2 = 29 at x=3"
print("✅ exact gradient, zero calculus homework")🎯 Two passes, not 130
On Day 4 you estimated each gradient by nudging one weight, running the whole network forward, and measuring the change — one forward pass per weight. With ~130 weights that was ~130 passes for a single training step.
loss.backward() just did that same job for all the weights at once, with one forward pass and one backward pass. That is the entire reason modern AI is possible: autograd records every operation on the way in, then replays the chain rule on the way out.
Under the hood it's all matrix multiplication — the same
@you just used. Making those multiplications scream on real hardware is a whole field of its own — high-performance computing.
3. The mission: read handwriting (MNIST)
MNIST is the "hello world" of computer vision: 60,000 training and 10,000 test handwritten digits (0–9), each a 28×28 grayscale image. 60,000 to train on, 10,000 strangers' digits to test on.
One rule matters for later: we apply only ToTensor(), which scales pixels to the range 0–1 and gives us white digits on a black background. We do not normalize — your Digit Reader Studio module expects exactly this format, so keep it in mind for the unlock.
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
tfm = transforms.ToTensor() # pixels -> [0, 1]; no normalization (module contract!)
train_full = datasets.MNIST(DATA_DIR, train=True, download=True, transform=tfm)
test_set = datasets.MNIST(DATA_DIR, train=False, download=True, transform=tfm)
train_set = Subset(train_full, range(TRAIN_SUBSET))
# num_workers=0: MNIST is tiny, so one process is plenty, and it avoids the
# harmless-but-noisy DataLoader worker-shutdown warnings you can hit in Jupyter.
train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
test_loader = DataLoader(test_set, batch_size=256, shuffle=False, num_workers=0)
print(f"train images: {len(train_set):,} test images: {len(test_set):,}")
img0, label0 = train_full[0]
print("one image:", tuple(img0.shape), "| pixel range:", float(img0.min()), "to", float(img0.max()))import matplotlib.pyplot as plt
fig, axes = plt.subplots(4, 8, figsize=(10, 5))
for ax, idx in zip(axes.ravel(), range(32)):
img, label = train_full[idx]
ax.imshow(img.squeeze(), cmap="gray")
ax.set_title(str(label), fontsize=9)
ax.axis("off")
fig.suptitle("MNIST — handwritten digits (white ink on black, values 0–1)")
plt.tight_layout(); plt.show()4. Build the network (a multi-layer perceptron)
Time to assemble the same kind of network you hand-wired on Day 4 — but in three lines. nn.Sequential just stacks layers front to back:
nn.Flatten()— unroll each 28×28 image into a flat vector of 784 numbersnn.Linear(784, 128)— the big matmul:(batch, 784) @ (784, 128)nn.ReLU()— the non-linearity that lets it learn curves, not just linesnn.Linear(128, 10)— ten output scores, one per digit
🧪 Exercise 3 — assemble the MLP
Build the four-layer nn.Sequential above and move it onto DEVICE.
import torch.nn as nn
# ==================== YOUR CODE HERE ====================
...
print(mlp)That nn.Linear(784, 128) is exactly the (batch, 784) @ (784, 128) matrix multiply you discovered on Day 4 — the network is made of matmuls. Let's count how many numbers it's juggling.
🧪 Exercise 4 — count the parameters
Sum p.numel() over mlp.parameters(). Expected: 101,770 (that's 784·128 + 128 weights and biases in the first layer, plus 128·10 + 10 in the second).
# ==================== YOUR CODE HERE ====================
### HINT: p.numel() counts the entries in one weight tensor; sum over all params
...
print(f"The MLP has {n_params:,} parameters.")
assert n_params == 101_770, "Expected 100,480 + 1,290 = 101,770 — check your layer sizes"
print("✅ 101,770 knobs, all tuned by autograd")5. The training loop — the pattern that runs all of AI
Almost every neural network in the world learns with the same five-step loop. Burn it into memory:
optimizer.zero_grad()— forget the last step's gradientslogits = model(x)— forward pass (a stack of matrix multiplies)loss = loss_fn(logits, y)— how wrong were we?loss.backward()— autograd fills in every gradientoptimizer.step()— nudge all the weights a little downhill
That's it. Even GPT-scale models train with this exact skeleton — just bigger.
🧪 Exercise 5 — complete the training loop
Fill the five ★ gaps in train_model below — one line each, in order. This is the single most important code you'll write all course.
- Loss function:
nn.CrossEntropyLoss()(already wired up). - Optimizer:
AdamatLEARNING_RATE(already wired up). - Expected: the printed loss goes down each epoch. If it stays flat or explodes, re-check the order of the five steps.
import torch.nn as nn
from tqdm.auto import tqdm
def train_model(model, epochs):
'Run the canonical training loop; reused for every model today.'
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
for epoch in range(1, epochs + 1):
model.train()
running = 0.0
for images, labels in tqdm(train_loader, desc=f"epoch {epoch}/{epochs}"):
images, labels = images.to(DEVICE), labels.to(DEVICE)
# ==================== YOUR CODE HERE ====================
### HINT: ★1 — clear the gradients left over from the previous step
...
# ==================== YOUR CODE HERE ====================
### HINT: ★2 — forward pass: feed the images through the model to get logits
...
# ==================== YOUR CODE HERE ====================
### HINT: ★3 — compare logits to the true labels to get the loss
...
# ==================== YOUR CODE HERE ====================
### HINT: ★4 — backpropagate: fill in a gradient for every weight
...
# ==================== YOUR CODE HERE ====================
### HINT: ★5 — take one step downhill with the optimizer
...
running += loss.item() * images.size(0)
print(f"epoch {epoch}: train loss {running / len(train_loader.dataset):.3f}")
return modeltrain_model(mlp, EPOCHS)
print("✅ MLP trained.")6. Grade it on strangers' handwriting
Training loss is nice, but the real question is: how well does it read digits it has never seen? We loop over the 10,000-image test set and count how many it gets right — all inside torch.no_grad(), because we're only reading the network now, not training it (no gradients needed, so it runs faster and uses less memory).
🧪 Exercise 6 — write evaluate
Fill in the no-grad loop that counts correct predictions. The class the network picks is the index of the largest logit: logits.argmax(dim=1).
def evaluate(model, loader):
'Return accuracy (fraction correct) over a data loader.'
model.eval()
correct = total = 0
# ==================== YOUR CODE HERE ====================
### HINT: wrap the loop in `with torch.no_grad():` and compare argmax to labels
...
return correct / total
mlp_acc = evaluate(mlp, test_loader)
print(f"MLP test accuracy: {mlp_acc:.4f}")
if not SMOKE:
assert mlp_acc >= 0.95, "Expected >= 0.95 on the full run — check your training loop"
print(f"🎉 Your network just read {round(mlp_acc * 10000):,} of 10,000 strangers' digits!")7. Why a GPU? Let's measure it ⚡
Everything above is matrix multiplication. A GPU does thousands of multiplies at once, which is why it obliterates a CPU on big matrices. Let's prove it: time A @ B for growing matrix sizes on the CPU, and (if you have one) on the GPU.
One catch: a GPU runs asynchronously. When you write A @ B, the CPU fires off the work and moves on before the GPU is done. So to time it honestly you must call torch.cuda.synchronize() — "wait until the GPU actually finishes" — before stopping the clock.
🧪 Exercise 7 — time the matmul
Fill in the two timed blocks (CPU and GPU). The warm-up call and the synchronize() calls are given — you write the timing with time.perf_counter().
import time
sizes = [256, 512, 1024] if SMOKE else [512, 1024, 2048, 4096]
cpu_times, gpu_times = [], []
for n in sizes:
A = torch.randn(n, n)
B = torch.randn(n, n)
# --- CPU timing ---
# ==================== YOUR CODE HERE ====================
### HINT: read time.perf_counter() before and after C = A @ B, store the difference
...
# --- GPU timing (only if we actually have a CUDA device) ---
if DEVICE.type == "cuda":
Ag, Bg = A.to(DEVICE), B.to(DEVICE)
_ = Ag @ Bg # warm-up: the very first CUDA call is slow
torch.cuda.synchronize() # wait for the warm-up to finish
# ==================== YOUR CODE HERE ====================
### HINT: same idea, but call torch.cuda.synchronize() before stopping the clock
...
line = f"n={n:5d} CPU {cpu_times[-1] * 1e3:8.1f} ms"
if DEVICE.type == "cuda":
line += f" GPU {gpu_times[-1] * 1e3:8.3f} ms"
print(line)import matplotlib.pyplot as plt
import numpy as np
x = np.arange(len(sizes))
plt.figure(figsize=(8, 5))
plt.bar(x - 0.2, np.array(cpu_times) * 1e3, width=0.4, label="CPU")
if DEVICE.type == "cuda":
plt.bar(x + 0.2, np.array(gpu_times) * 1e3, width=0.4, label="GPU")
plt.yscale("log")
plt.xticks(x, [str(n) for n in sizes])
plt.xlabel("matrix size n (an n × n matrix)")
plt.ylabel("time per matmul (ms, log scale)")
plt.title("Matrix multiply: CPU vs GPU")
plt.legend(); plt.tight_layout(); plt.show()
if DEVICE.type == "cuda":
for n, ct, gt in zip(sizes, cpu_times, gpu_times):
print(f"n={n:5d}: GPU is {ct / gt:6.1f}× faster")
else:
print("🐢 CPU only — no GPU to compare against right now.")
print(" On the Lab's GPU, the GPU bars would sit far below the CPU bars.")The numbers behind the numbers
Multiplying two n×n matrices costs about 2·n³ floating-point operations. Divide that by the time you measured and you get FLOP/s — operations per second. A workstation CPU manages a few billion (GFLOP/s); a modern GPU does tens of trillions (TFLOP/s). Same math, wildly different speed — because the GPU does thousands of multiplies in parallel.
Squeezing these numbers out of the hardware is a profession of its own — high-performance computing. Today you use the speed; one day you could build it.
8. A sharper tool: the convolutional network 🔬
The MLP flattens each image into 784 numbers on the very first line — throwing away the fact that nearby pixels belong together. A convolutional network keeps the 2-D shape and slides small pattern-detectors across it. (Day 6 goes deep on why this works — today, just feel the accuracy jump.)
Here's a small CNN. Don't worry about every line yet; notice it ends in the same two Linear layers, just fed by smarter features.
import torch.nn as nn
class SmallCNN(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 28 -> 14
nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 14 -> 7
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(32 * 7 * 7, 128), nn.ReLU(),
nn.Linear(128, 10),
)
def forward(self, x):
return self.classifier(self.features(x))
print(SmallCNN())🧪 Exercise 8 — train the CNN by reusing your loop
You already wrote train_model. That's the point of a reusable function: the CNN trains with the exact same call. Instantiate SmallCNN, move it to DEVICE, and train it.
cnn = SmallCNN().to(DEVICE)
# ==================== YOUR CODE HERE ====================
### HINT: no new loop needed — call the train_model you already wrote
...
cnn_acc = evaluate(cnn, test_loader)
print(f"CNN test accuracy: {cnn_acc:.4f}")print(f"MLP accuracy: {mlp_acc:.4f}")
print(f"CNN accuracy: {cnn_acc:.4f}")
print(f"The CNN got {(cnn_acc - mlp_acc) * 100:+.2f} percentage points by respecting the image's shape.")9. Autopsy: where does it slip?
No network is perfect. Let's look at the digits the CNN got wrong — true label on the left of the arrow, its guess on the right.
import matplotlib.pyplot as plt
cnn.eval()
wrong = []
with torch.no_grad():
for images, labels in test_loader:
preds = cnn(images.to(DEVICE)).argmax(1).cpu()
for img, t_, p_ in zip(images, labels, preds):
if int(t_) != int(p_):
wrong.append((img, int(t_), int(p_)))
if len(wrong) >= 16:
break
fig, axes = plt.subplots(2, 8, figsize=(10, 3))
for ax, (img, t_, p_) in zip(axes.ravel(), wrong[:16]):
ax.imshow(img.squeeze(), cmap="gray")
ax.set_title(f"{t_} → {p_}", fontsize=9, color="crimson")
ax.axis("off")
fig.suptitle("Where the network slips: true → predicted")
plt.tight_layout(); plt.show()Be honest with yourself as you look: how many of those would you get right, out of context? A lot of them are genuinely mangled — a 4 that's really a 9, a 7 with a slash that looks like a 1. The network isn't dumb; the handwriting is hard.
10. Read your handwriting ✍️
Now the fun part. The helper below turns whatever you scribble on the pad into the exact format your network trained on — a 28×28 white-on-black image, centered — then asks the CNN what it sees.
Tip: draw big and centered, filling most of the box. The network learned from centered digits, so a tiny scribble in the corner confuses it.
import numpy as np
from PIL import Image
def sketch_to_tensor(sketch):
'Sketchpad value -> MNIST-style (1,1,28,28) tensor + a preview, or (None, None).'
img = sketch.get("composite") if isinstance(sketch, dict) else sketch
if img is None:
return None, None
arr = np.asarray(img).astype(np.float32)
if arr.ndim == 3 and arr.shape[2] == 4: # RGBA -> flatten onto white
rgb, alpha = arr[..., :3], arr[..., 3:4] / 255.0
arr = rgb * alpha + 255.0 * (1.0 - alpha)
if arr.ndim == 3:
arr = arr.mean(axis=2) # to grayscale
ink = 255.0 - arr if np.median(arr) > 127 else arr # want white ink on black
if ink.max() < 10: # empty canvas
return None, None
ink = ink / ink.max() * 255.0
ys, xs = np.where(ink > 30) # crop tight to the strokes
crop = ink[ys.min():ys.max() + 1, xs.min():xs.max() + 1]
side = max(crop.shape)
sq = np.zeros((side, side), np.float32) # pad to a square
oy, ox = (side - crop.shape[0]) // 2, (side - crop.shape[1]) // 2
sq[oy:oy + crop.shape[0], ox:ox + crop.shape[1]] = crop
small = Image.fromarray(sq.astype(np.uint8)).resize((20, 20), Image.LANCZOS)
canvas = np.zeros((28, 28), np.float32)
canvas[4:24, 4:24] = np.asarray(small, dtype=np.float32) # center in 28x28
x = torch.from_numpy(canvas / 255.0).reshape(1, 1, 28, 28)
preview = np.kron(canvas.astype(np.uint8), np.ones((6, 6), np.uint8)) # zoom for display
return x, previewimport gradio as gr
cnn.eval()
def read_digit(sketch):
x, preview = sketch_to_tensor(sketch)
if x is None:
return {}, None
dev = next(cnn.parameters()).device # follow the model onto whatever device it is on
with torch.no_grad():
probs = torch.softmax(cnn(x.to(dev))[0], dim=0)
return {str(i): float(p) for i, p in enumerate(probs)}, preview
with gr.Blocks() as demo:
gr.Markdown("## ✍️ Draw a digit (0–9) — big and centered!")
with gr.Row():
pad = gr.Sketchpad(label="Draw here", canvas_size=(280, 280))
with gr.Column():
out = gr.Label(num_top_classes=3, label="I think it's a ...")
seen = gr.Image(label="What the network sees (28×28)", height=180)
pad.change(read_digit, pad, [out, seen])
if not SMOKE:
demo.launch(share=IN_COLAB)🔓 Unlock your Studio module
Time to ship. We export the trained CNN as a TorchScript file — a self-contained, frozen copy the AI Studio can load without your training code. The input contract must match exactly what your Studio module expects: a (1, 1, 28, 28) tensor with values 0–1, white-on-black.
# Export the trained CNN so the Studio's Digit Reader tab can load it.
REQUIRED_FILES = ["digit_reader.pt"]
cnn.eval()
example = torch.zeros(1, 1, 28, 28) # input contract: (1,1,28,28), values 0–1
scripted = torch.jit.trace(cnn.cpu(), example) # trace on CPU; the Studio runs on CPU
scripted.save(str(MODELS_DIR / "digit_reader.pt"))
cnn.to(DEVICE) # put the live model back on its device
missing = [f for f in REQUIRED_FILES if not (MODELS_DIR / f).exists()]
assert not missing, f"Something didn't save: {missing}"
print("🔓 UNLOCKED: Digit Reader! Run the Studio to see it live:")
print(" python ai_studio/app.py (from the repo root)")🏁 Checkpoint
Look what you built today:
- ✅ Traded hand-rolled gradients for autograd — two passes instead of 130
- ✅ Trained an MLP to ~97% on MNIST, then a CNN to even higher
- ✅ Measured a GPU demolishing a CPU at matrix multiply
- ✅ Drew your own digit and watched your network read it
- 🔓 Unlocked Digit Reader in your AI Studio
You fed this network tiny 28×28 gray digits and it learned to read. But the world isn't gray, and it isn't 28×28.
Tomorrow — Day 6, "Machines That See": real color photos, ten kinds of objects, and convolutions used in earnest. Your network is about to grow eyes. 👁️