beastgitea2026/gpt_from_scratch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

c

Browse Source
This commit is contained in:
betterbeast2023
2026-06-10 19:45:12 +08:00
parent 1183d10849
commit a3ebb28ce0
12 changed files with 35469 additions and 491 deletions
Download Patch File Download Diff File Expand all files Collapse all files
05_wavenet.ipynb
+424
View File
@@ -0,0 +1,424 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "d24a5a8b",
"metadata": {},
"source": [
"# makemore — a WaveNet\n",
"\n",
"Go **deeper and hierarchical** — **concept -> code -> Your turn** each step.\n",
"\n",
"This is notebook **05**. The MLP in `02`/`03` squashed the whole context into **one** layer.\n",
"Here we widen the context to 8 letters and fuse them **gradually**, two at a time, in a\n",
"tree — the structure behind DeepMind's WaveNet (2016). We also wrap everything in tidy\n",
"`torch.nn`-style building blocks. Follows Karpathy's *\"makemore part 5: WaveNet\"*."
]
},
{
"cell_type": "markdown",
"id": "187db6ed",
"metadata": {},
"source": [
"## Prologue — in plain English\n",
"\n",
"Reading 8 letters by mashing them into one layer throws away structure. WaveNet's idea: fuse\n",
"information **slowly**.\n",
"\n",
"Real-life picture: a knockout tournament. 8 players (letters) play in pairs -> 4 winners ->\n",
"play in pairs -> 2 -> 1 champion. Each round combines neighbours into a richer summary. Our\n",
"network does the same: combine letters in pairs, then pairs-of-pairs, then again, until a\n",
"single vector represents the whole context.\n",
"\n",
"We also build small reusable layers (`Linear`, `BatchNorm1d`, `Tanh`, `Embedding`,\n",
"`FlattenConsecutive`, `Sequential`) that behave like PyTorch's own — so the model reads like a\n",
"clean list of steps."
]
},
{
"cell_type": "markdown",
"id": "6bf209d0",
"metadata": {},
"source": [
"### 0.1 Imports and data (context grows to 8)\n",
"\n",
"The only data change from earlier notebooks: `block_size = 8`, so the model sees up to 8\n",
"previous letters."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "88825091",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ.setdefault(\"KMP_DUPLICATE_LIB_OK\", \"TRUE\")\n",
"\n",
"import random\n",
"import torch\n",
"import torch.nn.functional as F\n",
"\n",
"try:\n",
" import matplotlib.pyplot as plt\n",
" HAS_PLT = True\n",
"except Exception:\n",
" HAS_PLT = False\n",
"\n",
"words = open(\"names.txt\").read().splitlines()\n",
"chars = sorted(list(set(\"\".join(words))))\n",
"stoi = {ch: i + 1 for i, ch in enumerate(chars)}\n",
"stoi[\".\"] = 0\n",
"itos = {i: ch for ch, i in stoi.items()}\n",
"vocab_size = len(itos)\n",
"block_size = 8\n",
"\n",
"def build_dataset(word_list):\n",
" X, Y = [], []\n",
" for w in word_list:\n",
" context = [0] * block_size\n",
" for ch in w + \".\":\n",
" ix = stoi[ch]\n",
" X.append(context); Y.append(ix)\n",
" context = context[1:] + [ix]\n",
" return torch.tensor(X), torch.tensor(Y)\n",
"\n",
"random.seed(42)\n",
"random.shuffle(words)\n",
"n1, n2 = int(0.8 * len(words)), int(0.9 * len(words))\n",
"Xtr, Ytr = build_dataset(words[:n1])\n",
"Xdev, Ydev = build_dataset(words[n1:n2])\n",
"Xte, Yte = build_dataset(words[n2:])\n",
"print(\"train:\", tuple(Xtr.shape), \" (context length =\", block_size, \")\")"
]
},
{
"cell_type": "markdown",
"id": "4c8200dd",
"metadata": {},
"source": [
"### 1.1 PyTorch-style building blocks\n",
"\n",
"Each layer is a tiny class with a `__call__` (forward) and `parameters()`. The interesting one\n",
"is **`FlattenConsecutive(n)`**: it groups every `n` neighbouring positions together — that is\n",
"the \"combine in pairs\" tournament step. `BatchNorm1d` is written to also handle the 3-D tensors\n",
"that appear between fusion steps (a subtle but important detail)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7695c863",
"metadata": {},
"outputs": [],
"source": [
"class Linear:\n",
" def __init__(self, fan_in, fan_out, bias=True):\n",
" self.weight = torch.randn((fan_in, fan_out), generator=g) / fan_in ** 0.5\n",
" self.bias = torch.zeros(fan_out) if bias else None\n",
" def __call__(self, x):\n",
" self.out = x @ self.weight\n",
" if self.bias is not None:\n",
" self.out = self.out + self.bias\n",
" return self.out\n",
" def parameters(self):\n",
" return [self.weight] + ([] if self.bias is None else [self.bias])\n",
"\n",
"\n",
"class BatchNorm1d:\n",
" def __init__(self, dim, eps=1e-5, momentum=0.1):\n",
" self.eps = eps; self.momentum = momentum; self.training = True\n",
" self.gamma = torch.ones(dim); self.beta = torch.zeros(dim)\n",
" self.running_mean = torch.zeros(dim); self.running_var = torch.ones(dim)\n",
" def __call__(self, x):\n",
" if self.training:\n",
" dim = 0 if x.ndim == 2 else (0, 1) # 3-D inputs: average over batch AND time\n",
" xmean = x.mean(dim, keepdim=True)\n",
" xvar = x.var(dim, keepdim=True)\n",
" else:\n",
" xmean = self.running_mean\n",
" xvar = self.running_var\n",
" xhat = (x - xmean) / torch.sqrt(xvar + self.eps)\n",
" self.out = self.gamma * xhat + self.beta\n",
" if self.training:\n",
" with torch.no_grad():\n",
" self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean\n",
" self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar\n",
" return self.out\n",
" def parameters(self):\n",
" return [self.gamma, self.beta]\n",
"\n",
"\n",
"class Tanh:\n",
" def __call__(self, x):\n",
" self.out = torch.tanh(x); return self.out\n",
" def parameters(self):\n",
" return []\n",
"\n",
"\n",
"class Embedding:\n",
" def __init__(self, num_embeddings, embedding_dim):\n",
" self.weight = torch.randn((num_embeddings, embedding_dim), generator=g)\n",
" def __call__(self, IX):\n",
" self.out = self.weight[IX]; return self.out\n",
" def parameters(self):\n",
" return [self.weight]\n",
"\n",
"\n",
"class FlattenConsecutive:\n",
" def __init__(self, n):\n",
" self.n = n\n",
" def __call__(self, x):\n",
" B, T, C = x.shape\n",
" x = x.view(B, T // self.n, C * self.n) # fuse every n neighbours\n",
" if x.shape[1] == 1:\n",
" x = x.squeeze(1)\n",
" self.out = x; return self.out\n",
" def parameters(self):\n",
" return []\n",
"\n",
"\n",
"class Sequential:\n",
" def __init__(self, layers):\n",
" self.layers = layers\n",
" def __call__(self, x):\n",
" for layer in self.layers:\n",
" x = layer(x)\n",
" self.out = x; return self.out\n",
" def parameters(self):\n",
" return [p for layer in self.layers for p in layer.parameters()]\n",
"\n",
"print(\"layers defined\")"
]
},
{
"cell_type": "markdown",
"id": "ab8cd26d",
"metadata": {},
"source": [
"### 1.2 Build the hierarchical (WaveNet) model\n",
"\n",
"`block_size = 8` fuses in three rounds: 8 -> 4 -> 2 -> 1. Each round is\n",
"`FlattenConsecutive(2) -> Linear -> BatchNorm1d -> Tanh`. Finally a `Linear` maps to the 27\n",
"next-letter scores. We shrink the last layer's weights so the first loss starts near the ideal\n",
"`ln(27)`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1dcd57ee",
"metadata": {},
"outputs": [],
"source": [
"g = torch.Generator().manual_seed(42)\n",
"n_embd = 24\n",
"n_hidden = 128\n",
"\n",
"model = Sequential([\n",
" Embedding(vocab_size, n_embd),\n",
" FlattenConsecutive(2), Linear(n_embd * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),\n",
" FlattenConsecutive(2), Linear(n_hidden * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),\n",
" FlattenConsecutive(2), Linear(n_hidden * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),\n",
" Linear(n_hidden, vocab_size),\n",
"])\n",
"\n",
"with torch.no_grad():\n",
" model.layers[-1].weight *= 0.1 # humble output\n",
"\n",
"parameters = model.parameters()\n",
"for p in parameters:\n",
" p.requires_grad = True\n",
"print(\"total parameters:\", sum(p.nelement() for p in parameters))"
]
},
{
"cell_type": "markdown",
"id": "dfe5653b",
"metadata": {},
"source": [
"### 1.3 Watch the tournament shrink the sequence\n",
"\n",
"Run one batch and print each layer's output shape. See the time dimension halve 8 -> 4 -> 2 -> 1\n",
"as neighbours fuse, while the feature dimension does the work."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5a37517c",
"metadata": {},
"outputs": [],
"source": [
"ix = torch.randint(0, Xtr.shape[0], (4,), generator=g)\n",
"_ = model(Xtr[ix])\n",
"for layer in model.layers:\n",
" print(f\"{layer.__class__.__name__:18s} : {tuple(layer.out.shape)}\")"
]
},
{
"cell_type": "markdown",
"id": "43949ea9",
"metadata": {},
"source": [
"### 2.1 Train\n",
"\n",
"Standard loop: minibatch -> forward -> `loss.backward()` -> nudge. (We use autograd again now;\n",
"notebook `04` already proved we could do it by hand.)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f0800ce0",
"metadata": {},
"outputs": [],
"source": [
"max_steps = 20000\n",
"batch_size = 32\n",
"lossi = []\n",
"\n",
"for i in range(max_steps):\n",
" ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)\n",
" Xb, Yb = Xtr[ix], Ytr[ix]\n",
"\n",
" logits = model(Xb)\n",
" loss = F.cross_entropy(logits, Yb)\n",
"\n",
" for p in parameters:\n",
" p.grad = None\n",
" loss.backward()\n",
" lr = 0.1 if i < 15000 else 0.01\n",
" for p in parameters:\n",
" p.data += -lr * p.grad\n",
"\n",
" lossi.append(loss.log10().item())\n",
" if i % 4000 == 0 or i == max_steps - 1:\n",
" print(f\"step {i:5d}/{max_steps} loss {loss.item():.4f}\")\n",
"\n",
"if HAS_PLT:\n",
" # smooth the noisy per-step loss by averaging in chunks of 200\n",
" lt = torch.tensor(lossi)\n",
" plt.figure(figsize=(5, 4))\n",
" plt.plot(lt.view(-1, 200).mean(1))\n",
" plt.xlabel(\"step / 200\"); plt.ylabel(\"log10(loss)\"); plt.title(\"WaveNet training loss\")\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "053ba8a7",
"metadata": {},
"source": [
"### 2.2 Evaluate\n",
"\n",
"Switch BatchNorm layers to eval mode (use running stats), then score train and dev. The deeper\n",
"hierarchical model should edge below the single-layer MLP's ~2.1."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3a5d6228",
"metadata": {},
"outputs": [],
"source": [
"for layer in model.layers:\n",
" if isinstance(layer, BatchNorm1d):\n",
" layer.training = False\n",
"\n",
"@torch.no_grad()\n",
"def split_loss(split):\n",
" x, y = {\"train\": (Xtr, Ytr), \"dev\": (Xdev, Ydev), \"test\": (Xte, Yte)}[split]\n",
" logits = model(x)\n",
" return F.cross_entropy(logits, y).item()\n",
"\n",
"print(\"train loss:\", round(split_loss(\"train\"), 4))\n",
"print(\"dev loss:\", round(split_loss(\"dev\"), 4))"
]
},
{
"cell_type": "markdown",
"id": "b78eaf70",
"metadata": {},
"source": [
"### 3.1 Generate names\n",
"\n",
"Sample letter by letter as before. With 8 letters of context and the hierarchical fusion, the\n",
"names tend to look a little more name-like."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "33833066",
"metadata": {},
"outputs": [],
"source": [
"g2 = torch.Generator().manual_seed(2147483647 + 10)\n",
"\n",
"for _ in range(20):\n",
" out = []\n",
" context = [0] * block_size\n",
" while True:\n",
" logits = model(torch.tensor([context]))\n",
" probs = F.softmax(logits, dim=1)\n",
" ix = torch.multinomial(probs, num_samples=1, generator=g2).item()\n",
" context = context[1:] + [ix]\n",
" if ix == 0:\n",
" break\n",
" out.append(itos[ix])\n",
" print(\"\".join(out))"
]
},
{
"cell_type": "markdown",
"id": "31d52503",
"metadata": {},
"source": [
"**Your turn 5** — The model is just a list. Add another fusion round (so it handles\n",
"`block_size = 16`), or raise `n_hidden`, retrain, and compare dev loss. Because everything is a\n",
"clean `Sequential`, experimenting is just editing the list."
]
},
{
"cell_type": "markdown",
"id": "894cc92e",
"metadata": {},
"source": [
"## What's next\n",
"\n",
"You built a deeper, hierarchical model out of reusable PyTorch-style blocks — the same shape of\n",
"code real libraries use.\n",
"\n",
"Where this leads in the series:\n",
"\n",
"- `06_build_gpt_attention.ipynb` — the transformer. Instead of a fixed tournament that always\n",
" fuses neighbours, **attention** lets every position decide *which* earlier positions to listen\n",
" to. That flexibility is what makes GPT powerful.\n",
"- `07_gpt_tokenizer.ipynb` — how real models turn text into tokens in the first place.\n",
"\n",
"**Checklist**\n",
"- [ ] Context widened to 8 letters\n",
"- [ ] PyTorch-style layers (Linear, BatchNorm1d, Tanh, Embedding, FlattenConsecutive, Sequential)\n",
"- [ ] Hierarchical fusion 8 -> 4 -> 2 -> 1\n",
"- [ ] BatchNorm1d handles 3-D tensors\n",
"- [ ] Trained, evaluated, sampled names"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv (3.12.10)",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}