{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Lab 8 — Adding attention to the digit reader\n", "\n", "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/kleinric/cv-labs/blob/main/lab-08.ipynb)\n", "\n", "**COMS4036A / COMS7050A Computer Vision**\n", "\n", "Groups of up to three; one member submits the group's notebook (`.ipynb`) on Moodle. Every member must be able to explain every cell.\n", "\n", "Companion reading is [Chapter 8 of the course book](https://courses.ms.wits.ac.za/~richard/cv/book/chapters/08-attention-transformers.html). The digit reader from Lab 7 starts its decoder from one summary of the strip, $s_0$, and never looks at the image again. In this lab you give the reader a state for every column of the CNN's feature map, build two readers that differ mainly in whether the decoder can return to those states, and compare them on strips of six to eight digits and on strips longer than any they were trained on. After twelve epochs on strips of one to eight digits, the reference fixed-context reader reads about 77% of the test strips of six to eight digits exactly and the attention reader about 96%; three reference runs, with seeds 0, 1 and 2, gave 0.763, 0.771 and 0.783 for the first and 0.969, 0.961 and 0.962 for the second. Your run will differ in detail, and the behaviour on strips longer than the training lengths varies from run to run more than these numbers do.\n", "\n", "This lab's rung of the skills ladder is **a controlled comparison: change as little as possible, and say what else changed**. The two readers you build use the same encoder design, data, training recipe and seeds, and differ mainly in whether the decoder can return to the column states. Adding attention also adds parameters: the attention reader is about 16% larger, because of its projections and its wider decoder input and output layers. The comparison therefore tests this change to the architecture without separating attention from the extra capacity, and you report both parameter counts so that a reader can take that into account. The same care applies to the attention maps. A weight is a coefficient in the mixture that forms a context; it does not by itself show which part of the image a prediction depended on. Section 5 tests that by replacing a digit and measuring how the readings change, rather than by looking at a map.\n", "\n", "## 0. GPU, W&B, and the data\n", "\n", "Open this notebook in [Colab](https://colab.research.google.com) with the badge above, or [download it from the book](https://courses.ms.wits.ac.za/~richard/cv/book/labs/lab-08.ipynb) and use **File ▸ Upload notebook**; then **File ▸ Save a copy in Drive**. Set **Runtime ▸ Change runtime type** to a GPU, and log in to W&B; this lab logs to a project called `cv-lab8`." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# Group members — fill in before submitting.\n", "MEMBERS = [\n", " # (\"Student name\", \"Student number\"),\n", "]\n", "for name, number in MEMBERS:\n", " print(f\"{number} {name}\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "!nvidia-smi --query-gpu=name,memory.total --format=csv" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "import wandb\n", "wandb.login()" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "import random\n", "\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from torch.nn.utils.rnn import pad_sequence\n", "from torch.utils.data import DataLoader, Dataset\n", "from torchvision import datasets\n", "\n", "DEV = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "\n", "\n", "def seed_everything(seed=0):\n", " random.seed(seed)\n", " np.random.seed(seed)\n", " torch.manual_seed(seed)\n", " torch.cuda.manual_seed_all(seed)\n", "\n", "\n", "torch.backends.cudnn.benchmark = False # repeatable runs on the same GPU\n", "torch.backends.cudnn.deterministic = True\n", "seed_everything(0)\n", "print(\"device:\", DEV)" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The digits are MNIST, as in Lab 7: 60,000 training and 10,000 test images of single handwritten digits, each $28 \\times 28$ pixels (LeCun, Cortes and Burges). Training strips are built only from the training digits and test strips only from the test digits, so no handwritten digit in a test strip has been seen in training." ] }, { "cell_type": "code", "metadata": {}, "source": [ "mnist_train = datasets.MNIST(\"data\", train=True, download=True)\n", "mnist_test = datasets.MNIST(\"data\", train=False, download=True)\n", "train_imgs, train_labels = mnist_train.data.numpy(), mnist_train.targets.numpy()\n", "test_imgs, test_labels = mnist_test.data.numpy(), mnist_test.targets.numpy()\n", "print(f\"MNIST digits: {len(train_imgs)} train, {len(test_imgs)} test\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The vocabulary and the scoring functions are Lab 7's. The digits 0 to 9 are their own indices, and 10, 11 and 12 are ``, `` and ``. `clean` turns a token sequence into a list of digits by stopping at the first `` and dropping `` and ``. **Exact-string accuracy** is the share of strips whose cleaned prediction equals the target. **Character error rate** (CER) is the number of digit substitutions, deletions and insertions needed to turn the prediction into the target, summed over the strips and divided by the total number of target digits. `by_length` reports both measures separately for each target length." ] }, { "cell_type": "code", "metadata": {}, "source": [ "SOS, EOS, PAD = 10, 11, 12\n", "VOCAB = 13\n", "SYMBOLS = [str(d) for d in range(10)] + [\"\", \"\", \"\"]\n", "\n", "\n", "def show_tokens(seq):\n", " return \" \".join(SYMBOLS[int(t)] for t in seq)\n", "\n", "\n", "def clean(seq):\n", " \"\"\"Digits only: stop at the first , drop and .\"\"\"\n", " out = []\n", " for t in seq:\n", " t = int(t)\n", " if t == EOS:\n", " break\n", " if t < 10:\n", " out.append(t)\n", " return out\n", "\n", "\n", "def edit_distance(a, b):\n", " \"\"\"Minimum substitutions, deletions and insertions turning a into b.\"\"\"\n", " prev = list(range(len(b) + 1))\n", " for i in range(1, len(a) + 1):\n", " cur = [i] + [0] * len(b)\n", " for j in range(1, len(b) + 1):\n", " cur[j] = min(prev[j] + 1, cur[j - 1] + 1,\n", " prev[j - 1] + (a[i - 1] != b[j - 1]))\n", " prev = cur\n", " return prev[-1]\n", "\n", "\n", "def exact_string_accuracy(preds, targets):\n", " \"\"\"Fraction of strings whose cleaned prediction equals the cleaned target.\"\"\"\n", " return float(np.mean([clean(p) == clean(t) for p, t in zip(preds, targets)]))\n", "\n", "\n", "def cer(preds, targets):\n", " \"\"\"Character error rate: total edits over total target digits.\"\"\"\n", " edits = sum(edit_distance(clean(p), clean(t)) for p, t in zip(preds, targets))\n", " return edits / sum(len(clean(t)) for t in targets)\n", "\n", "\n", "def by_length(preds, targets):\n", " \"\"\"{length: (exact-string accuracy, CER, number of strips)}, by target length.\"\"\"\n", " rows = {}\n", " for p, t in zip(preds, targets):\n", " rows.setdefault(len(clean(t)), []).append((p, t))\n", " return {L: (exact_string_accuracy(*zip(*r)), cer(*zip(*r)), len(r))\n", " for L, r in sorted(rows.items())}" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Strips on a wider canvas\n", "\n", "The generator is Lab 7's, with the width of the canvas as a parameter. Lab 7's canvas was 224 pixels wide, enough for eight digits. Here it is 336 pixels wide, so that a strip of twelve digits always fits: twelve 28-pixel boxes with no overlap take exactly 336 pixels. As before, neighbouring boxes overlap by between 0 and 6 pixels and the whole string starts at a random horizontal position. The generator now also keeps, for every strip, the index of each digit's source image (`dataset.src[i]`) and the left edge of each digit's box (`dataset.xs[i]`), so that Section 5 can compare the attention weights with where the digits are, and redraw a strip with one digit changed using `draw_strip`." ] }, { "cell_type": "code", "metadata": {}, "source": [ "H, W = 28, 336\n", "GAP_RANGE = (-6, 0)\n", "\n", "\n", "def draw_strip(images, idx, xs, width=W):\n", " \"\"\"Draw the digits images[idx[k]] with their left edges at xs[k].\"\"\"\n", " canvas = np.zeros((H, width), dtype=np.uint8)\n", " for i, x in zip(idx, xs):\n", " np.maximum(canvas[:, x:x + 28], images[i], out=canvas[:, x:x + 28])\n", " return canvas\n", "\n", "\n", "def make_strip(images, labels, length, rng, width=W):\n", " \"\"\"One strip of `length` digits on a blank canvas 28 pixels high and `width` wide.\n", "\n", " Lab 7's generator, with the canvas width as a parameter. Returns the canvas\n", " (uint8), the digits, the indices of the source images and the left edge of\n", " each digit's 28-pixel box.\n", " \"\"\"\n", " idx = rng.integers(0, len(images), size=length)\n", " gaps = rng.integers(GAP_RANGE[0], GAP_RANGE[1] + 1, size=length - 1)\n", " span = 28 * length + int(gaps.sum())\n", " assert span <= width, f\"{length} digits do not fit on a canvas {width} wide\"\n", " x0 = int(rng.integers(0, width - span + 1))\n", " xs = x0 + np.concatenate([[0], np.cumsum(28 + gaps)]).astype(int)\n", " return draw_strip(images, idx, xs, width), [int(labels[i]) for i in idx], idx, xs\n", "\n", "\n", "class StripDataset(Dataset):\n", " \"\"\"`n` strips, equally many of each length in `lengths`, fixed by `seed`.\n", "\n", " Item i is (image, digits): a float tensor of shape (1, 28, width) in [0, 1]\n", " and a long tensor holding the digits only. The source-image indices and the\n", " box left edges of strip i are in self.src[i] and self.xs[i].\n", " \"\"\"\n", "\n", " def __init__(self, images, labels, n, lengths, seed, width=W):\n", " rng = np.random.default_rng(seed)\n", " lens = np.resize(np.asarray(lengths), n) # equal counts per length\n", " rng.shuffle(lens)\n", " self.images = np.zeros((n, H, width), dtype=np.uint8)\n", " self.digits, self.src, self.xs = [], [], []\n", " for j, L in enumerate(lens):\n", " self.images[j], d, idx, xs = make_strip(images, labels, int(L), rng, width)\n", " self.digits.append(d)\n", " self.src.append(idx)\n", " self.xs.append(xs)\n", "\n", " def __len__(self):\n", " return len(self.digits)\n", "\n", " def __getitem__(self, i):\n", " img = torch.from_numpy(self.images[i]).float().div(255).unsqueeze(0)\n", " return img, torch.tensor(self.digits[i], dtype=torch.long)\n", "\n", "\n", "test_set = StripDataset(test_imgs, test_labels, 8_000, range(1, 9), seed=1234)\n", "long_set = StripDataset(test_imgs, test_labels, 4_000, range(9, 13), seed=4321)\n", "print(len(test_set), len(long_set))" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "`test_set` holds 1,000 strips of each length from one to eight, and `long_set` 1,000 of each length from nine to twelve, all from MNIST test digits and identical on every run. Both readers train on strips of one to eight digits, so neither trains on strips of nine to twelve digits; Section 6 uses `long_set`.\n", "\n", "The batching function and the masked loss are the ones you wrote in Lab 7. `collate` pads the decoder inputs (`` followed by the digits) and the targets (the digits followed by ``) with `` to a common length, and `masked_loss` averages the cross-entropy over every target position that is not ``. The CNN is Lab 7's too: four blocks, each a $3 \\times 3$ convolution, batch norm, ReLU and a $2 \\times 2$ max-pool, at 32, 64, 128 and 128 channels." ] }, { "cell_type": "code", "metadata": {}, "source": [ "def collate(batch):\n", " \"\"\"Stack the images; pad decoder inputs and targets with PAD.\"\"\"\n", " images = torch.stack([img for img, _ in batch])\n", " sos, eos = torch.tensor([SOS]), torch.tensor([EOS])\n", " dec_in = pad_sequence([torch.cat([sos, d]) for _, d in batch],\n", " batch_first=True, padding_value=PAD)\n", " tgt = pad_sequence([torch.cat([d, eos]) for _, d in batch],\n", " batch_first=True, padding_value=PAD)\n", " return images, dec_in, tgt\n", "\n", "\n", "def masked_loss(logits, tgt):\n", " \"\"\"Mean cross-entropy over the target positions that are not .\"\"\"\n", " return F.cross_entropy(logits.reshape(-1, VOCAB), tgt.reshape(-1), ignore_index=PAD)\n", "\n", "\n", "def conv_block(c_in, c_out):\n", " return nn.Sequential(nn.Conv2d(c_in, c_out, 3, padding=1, bias=False),\n", " nn.BatchNorm2d(c_out), nn.ReLU(), nn.MaxPool2d(2))\n", "\n", "\n", "def make_cnn():\n", " \"\"\"Lab 7's CNN: four blocks at 32, 64, 128 and 128 channels.\"\"\"\n", " return nn.Sequential(conv_block(1, 32), conv_block(32, 64),\n", " conv_block(64, 128), conv_block(128, 128))" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Q1.1.** Display six strips from `test_set` and four from `long_set`, each titled with its digits and with each digit's box outlined from `dataset.xs`.\n", "\n", "Before running the CNN, work out three numbers: the shape of the CNN's feature map $F$ for one strip; the distance in pixels between the strip positions that neighbouring columns of $F$ correspond to; and how many columns a 28-pixel digit box spans at that spacing. Check the shape with `make_cnn()`.\n", "\n", "Finally, find the window of input pixels that one column of $F$ depends on, its *receptive field*: compute it from the layers by hand, and check it by back-propagating from one column of $F$ of an untrained CNN in evaluation mode, given a random image (`torch.rand`) rather than a blank one, and finding which pixel columns of the input receive a non-zero gradient. Why would a blank image not work for this check? How many digit boxes can that window overlap?" ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "## 2. A state for every column\n", "\n", "The CNN turns the strip into 21 columns of 128 numbers each. In Lab 7 the decoder's first state was a projection of all of them at once, and the decoder never saw them again. Attention needs the opposite: a collection of states, one per column, that the decoder can return to at every step.\n", "\n", "A column of $F$ describes only its own window of the strip. A **bidirectional GRU** makes each column's state describe that column in the context of the whole strip. One GRU reads the 21 columns from left to right and another reads them from right to left, and the state for column $j$ joins the two:\n", "\n", "$$\n", "h_j=\\bigl[\\overrightarrow{h}_j;\\,\\overleftarrow{h}_j\\bigr],\\qquad j=1,\\ldots,21.\n", "$$\n", "\n", "The forward state $\\overrightarrow{h}_j$ has read columns 1 to $j$, and the backward state $\\overleftarrow{h}_j$ has read columns 21 down to $j$. With 64 units in each direction, each $h_j$ has 128 numbers. The decoder's first state comes from the two final states, the forward one after column 21 and the backward one after column 1,\n", "\n", "$$\n", "s_0=\\tanh\\!\\left(W_{\\text{init}}\\bigl[\\overrightarrow{h}_{21};\\,\\overleftarrow{h}_1\\bigr]+b_{\\text{init}}\\right),\n", "$$\n", "\n", "with a width of 64. Both readers in Section 3 use this encoder, so they start from the same kind of $s_0$ and differ in what the decoder can read after that.\n", "\n", "`nn.GRU(..., bidirectional=True, batch_first=True)` does both directions in one call. Its first output has shape $(B, 21, 128)$, with the forward states in the first 64 numbers of each column and the backward states in the last 64. Its second output has shape $(2, B, 64)$: index 0 is the forward GRU's final state and index 1 the backward GRU's.\n", "\n", "**Q2.1.** Implement `ColumnEncoder(hidden=64, units=64)` with three parts named `cnn` (from `make_cnn()`), `bigru` (an `nn.GRU` from 128 inputs to `units` units in each direction) and `init` (an `nn.Linear(2 * units, hidden)`). Its `forward(images, trace=False)` turns $F$ of shape $(B, 128, 1, 21)$ into a sequence of 21 column vectors, runs the bidirectional GRU over them, and returns two tensors: `memory`, the column states $h_1, \\ldots, h_{21}$ as $(B, 21, 128)$, and $s_0$ as $(B, 64)$. When `trace` is true, print the shapes of $F$, the column sequence, `memory`, the final states and $s_0$." ] }, { "cell_type": "code", "metadata": {}, "source": [ "class ColumnEncoder(nn.Module):\n", " # YOUR CODE HERE\n", " ..." ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Q2.2.** Before you run it, compute by hand the number of parameters in the bidirectional GRU and in $W_{\\text{init}}$ and $b_{\\text{init}}$. A GRU layer with $n$ inputs and $u$ units holds three sets of input weights ($u \\times n$), recurrent weights ($u \\times u$) and two biases ($u$ each). Then call `seed_everything(0)`, build a `ColumnEncoder`, run the trace on a batch of four strips from `test_set`, and print the parameter count of each part.\n", "\n", "Then test the claim that a column state describes the whole strip. Take the first eight-digit strip in `test_set`, and draw a second version of it with its last digit erased (`draw_strip` with the last entry of `src` and `xs` left out). Put the encoder in evaluation mode, so that batch norm treats the two strips independently, and compute `memory` for both. Predict, then report, the largest absolute change in the forward half and in the backward half of the states of column 1 and of column 21, and list the columns whose forward half changes at all. Explain the pattern from the directions in which the two GRUs read and from the receptive field you found in Q1.1." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "## 3. Two readers\n", "\n", "Both readers have the same three methods, which the supplied training and generation code in Section 4 calls:\n", "\n", "- `encode(images)` returns $s_0$, of shape $(B, 64)$, and `memory`, whatever the decoder keeps from the image for the steps that follow;\n", "- `step(y_prev, s_prev, memory)` runs one decoder step from the previous symbol, of shape $(B,)$, and the previous state, and returns the logits $(B, 13)$, the new state $(B, 64)$, and the attention weights $(B, 21)$, or `None` for a reader without attention;\n", "- `forward(images, dec_in)` runs the decoder under teacher forcing over every position of `dec_in` and returns logits of shape $(B, T, 13)$.\n", "\n", "The **fixed-context reader** is supplied. Its decoder is Lab 7's: $s_t = \\operatorname{GRU}(E[y_{t-1}], s_{t-1})$ and $\\ell_t = W_{\\text{out}} s_t + b_{\\text{out}}$, with an embedding width of 32, starting from the $s_0$ of the column encoder. Its `encode` returns `None` as the memory. The column states are computed, but they reach the decoder only through $s_0$, and everything the decoder learns about this particular strip must pass through those 64 numbers." ] }, { "cell_type": "code", "metadata": {}, "source": [ "class FixedReader(nn.Module):\n", " \"\"\"The column encoder's s_0, then Lab 7's GRU decoder. The decoder never\n", " reads the column states: encode() returns None as the memory.\"\"\"\n", "\n", " def __init__(self, hidden=64, emb=32):\n", " super().__init__()\n", " self.encoder = ColumnEncoder(hidden)\n", " self.emb = nn.Embedding(VOCAB, emb)\n", " self.gru = nn.GRU(emb, hidden, batch_first=True)\n", " self.out = nn.Linear(hidden, VOCAB)\n", "\n", " def encode(self, images):\n", " memory, s0 = self.encoder(images)\n", " return s0, None\n", "\n", " def step(self, y_prev, s_prev, memory):\n", " o, s = self.gru(self.emb(y_prev).unsqueeze(1), s_prev.unsqueeze(0))\n", " return self.out(o[:, 0]), s[0], None\n", "\n", " def forward(self, images, dec_in):\n", " s0, _ = self.encode(images)\n", " o, _ = self.gru(self.emb(dec_in), s0.unsqueeze(0))\n", " return self.out(o)" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The **attention reader** keeps the column states and reads them again at every output step with scaled dot-product attention. The preceding decoder state supplies the query and each column state supplies a key and a value,\n", "\n", "$$\n", "q_t=W_Q\\,s_{t-1},\\qquad k_j=W_K\\,h_j,\\qquad v_j=W_V\\,h_j,\n", "$$\n", "\n", "with $d_k = 64$ for the query and the keys and $d_v = 128$ for the values, and no biases. These are written for one vector at a time, as column vectors, which is how Lab 7 wrote $W_{\\text{init}}$. The chapter stacks vectors as the rows of a matrix and writes $Q = X_q W_Q$, so its $W_Q$ is the transpose of the one here: $d_q \\times d_k$ in the chapter, $d_k \\times d_q$ here. `nn.Linear(64, 64, bias=False)` stores its weight in the shape used here, (out, in), and applied to a batch of row vectors it computes $X W^\\top$, so the two conventions give the same numbers. The scores are normalised by a softmax over the 21 columns, and the context is the weighted sum of the values:\n", "\n", "$$\n", "e_{tj}=\\frac{q_t\\cdot k_j}{\\sqrt{d_k}},\\qquad\n", "\\alpha_{tj}=\\frac{\\exp(e_{tj})}{\\sum_{i=1}^{21}\\exp(e_{ti})},\\qquad\n", "c_t=\\sum_{j=1}^{21}\\alpha_{tj}\\,v_j.\n", "$$\n", "\n", "The decoder then takes the context in two places, as part of the GRU's input and as part of the output layer's input:\n", "\n", "$$\n", "s_t=\\operatorname{GRU}\\bigl([E[y_{t-1}];\\,c_t],\\,s_{t-1}\\bigr),\\qquad\n", "\\ell_t=W_{\\text{out}}[s_t;\\,c_t]+b_{\\text{out}}.\n", "$$\n", "\n", "The keys and values depend only on the image, so they are computed once per strip. The same $W_Q$, $W_K$, $W_V$ and decoder parameters are used at every step; what changes from step to step is the query, and so the weights.\n", "\n", "**Q3.1.** Write `attend(q, K, V)`, scaled dot-product attention for one query per row of a batch: `q` has shape $(B, d_k)$, `K` $(B, N, d_k)$ and `V` $(B, N, d_v)$, and it returns the context $(B, d_v)$ and the weights $(B, N)$. Do not call `F.scaled_dot_product_attention` inside it. Then check it in three ways.\n", "\n", "1. Use the chapter's worked example from the attention mixer, a single query $q = (1.2, 0.6)$, keys $(1.4, 0.9)$, $(-1.1, 1.3)$ and $(-0.8, -1.4)$, and values $(0.86, 0.20, 0.18)$, $(0.22, 0.62, 0.28)$ and $(0.48, 0.30, 0.78)$, and compare your weights and output with the chapter's.\n", "2. Use a zero query. What must the weights be, whatever the keys?\n", "3. Use random tensors of the sizes this lab uses ($B = 5$, $N = 21$, $d_k = 64$, $d_v = 128$, after `torch.manual_seed(0)`): check that every row of weights sums to one, and report the largest difference between your context and `F.scaled_dot_product_attention` given the query as a sequence of length one." ] }, { "cell_type": "code", "metadata": {}, "source": [ "def attend(q, K, V):\n", " # YOUR CODE HERE\n", " ..." ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "**Q3.2.** Implement `AttentionReader(hidden=64, emb=32, d_k=64, d_v=128)` with the three methods above. It needs a `ColumnEncoder(hidden)` named `encoder`, an embedding table $E$, the projections $W_Q$, $W_K$ and $W_V$ as `nn.Linear` layers without biases, an `nn.GRUCell` whose input is the embedding and the context joined, and the output layer $W_{\\text{out}}$ on the state and the context joined. `encode` returns $s_0$ and, as the memory, the keys and values of the 21 columns. `step` computes the query from `s_prev`, calls your `attend`, updates the state and returns the logits, the new state and the weights. `forward` calls `encode` once and then `step` at each position of `dec_in`." ] }, { "cell_type": "code", "metadata": {}, "source": [ "class AttentionReader(nn.Module):\n", " # YOUR CODE HERE\n", " ..." ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Q3.3.** Before you count them: which parameters does the attention reader have that the fixed-context reader does not, and which of the fixed-context reader's layers become larger, and by how much? Compute the difference by hand, then print the parameter count of each part of both readers, the two totals and their ratio. Then the ritual from Lab 2: build one batch of 32 training strips, `StripDataset(train_imgs, train_labels, 32, range(1, 9), seed=999)`, call `seed_everything(0)`, and drive a fresh `AttentionReader` to memorise it with Adam at $10^{-3}$ for 300 steps under teacher forcing with `masked_loss`. Log it to W&B under a name that makes clear it is the ritual, print the loss at step 1 and every 50 steps, and report the first step at which it falls below 0.01." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "## 4. Train and compare\n", "\n", "The generation and training loops are supplied, because the recipe is Lab 7's and the comparison depends on both readers following it exactly. `generate` decodes greedily from ``, feeding back the highest-scoring symbol, and stops once every strip in the batch has written `` or after 14 steps, which is enough for a twelve-digit string and its ``. It returns the symbols, $(B, \\text{steps})$, and for the attention reader the weights, $(B, \\text{steps}, 21)$. `predict` runs it over a dataset and returns lists of predictions, targets and, for the attention reader, one array of weights per strip. Both put the model in evaluation mode, track no gradients, and return the model to the mode it was in.\n", "\n", "The settings of the recipe are collected in one cell below. Change them only to check that your code runs end to end, and restore them before the runs you report. The one exception is `N_STRIPS`, which may instead be 30,000, as explained after the cell.\n", "\n", "`train_reader` follows Lab 7's recipe on strips of one to eight digits: a fresh set of 60,000 training strips each epoch, batches of 128, Adam at $10^{-3}$ with cosine annealing stepped after every batch, twelve epochs, teacher forcing and `masked_loss`. It logs the mean training loss after every epoch, and the free-running exact-string accuracy and CER on `test_set` after the epochs listed in `EVAL_EPOCHS`, which saves evaluation time without changing the training. As in Lab 7 there is no separate validation split, so the per-epoch test numbers are for illustration only: choose the training budget before running either reader, use the same budget for both, and do not use the test curves to choose any setting." ] }, { "cell_type": "code", "metadata": {}, "source": [ "EPOCHS = 12 # training epochs for each reader\n", "N_STRIPS = 60_000 # fresh training strips per epoch\n", "BATCH = 128 # training batch size\n", "EVAL_EPOCHS = (1, 3, 6, 9, 12) # epochs after which test_set is scored and logged\n", "EVAL_BATCH = 500 # batch size for generation" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Training the two readers takes nearly all of a reference run's 38 minutes on a Colab T4, and the training time is roughly proportional to `N_STRIPS`. If you cannot spare that, you may set `N_STRIPS = 30_000` for both readers, provided you say so in your answer to Q4.3; the W&B config records it as well. No other value is permitted. The cost of the smaller budget falls unevenly on the two readers. In three reference runs at 30,000 strips per epoch, the attention reader still read about 96% of the strips of six to eight digits exactly, but the fixed-context reader fell from about 77% to about 40%. The fixed-context reader is still improving when training stops, so the smaller budget widens the gap between the readers, and part of the wider gap then reflects how much data each reader needs rather than what attention adds. If you use 30,000, your comparison in Section 4 has to say so." ] }, { "cell_type": "code", "metadata": {}, "source": [ "@torch.no_grad()\n", "def generate(model, images, max_len=14):\n", " \"\"\"Greedy decoding from ; stops once every row has written .\n", "\n", " Returns the symbols (B, steps) and the attention weights (B, steps, 21),\n", " or None for a reader without attention.\n", " \"\"\"\n", " was_training = model.training\n", " model.eval()\n", " s, memory = model.encode(images)\n", " y = torch.full((images.shape[0],), SOS, dtype=torch.long, device=images.device)\n", " out, weights = [], []\n", " for _ in range(max_len):\n", " logits, s, alpha = model.step(y, s, memory)\n", " y = logits.argmax(-1)\n", " out.append(y)\n", " if alpha is not None:\n", " weights.append(alpha)\n", " if (torch.stack(out, 1) == EOS).any(1).all():\n", " break\n", " model.train(was_training)\n", " return torch.stack(out, 1), (torch.stack(weights, 1) if weights else None)\n", "\n", "\n", "@torch.no_grad()\n", "def predict(model, dataset, batch_size=EVAL_BATCH):\n", " \"\"\"Free-running predictions, targets and (attention only) weights, as lists.\"\"\"\n", " preds, targets, weights = [], [], []\n", " for images, _, tgt in DataLoader(dataset, batch_size=batch_size, collate_fn=collate):\n", " p, a = generate(model, images.to(DEV))\n", " preds += p.tolist()\n", " targets += tgt.tolist()\n", " if a is not None:\n", " weights += list(a.cpu().numpy())\n", " return preds, targets, (weights if weights else None)\n", "\n", "\n", "def train_reader(make_model, seed, name, lengths=range(1, 9)):\n", " \"\"\"Lab 7's recipe on strips of `lengths`, logged to W&B. Returns the model\n", " and a list of per-epoch rows.\"\"\"\n", " seed_everything(seed)\n", " model = make_model().to(DEV)\n", " opt = torch.optim.Adam(model.parameters(), lr=1e-3)\n", " steps_per_epoch = N_STRIPS // BATCH # 468\n", " sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, EPOCHS * steps_per_epoch)\n", " g = torch.Generator().manual_seed(seed)\n", " config = dict(reader=type(model).__name__, hidden=64, emb=32, bigru_units=64,\n", " train_lengths=list(lengths), strips_per_epoch=N_STRIPS,\n", " optimiser=\"Adam\", lr=1e-3, schedule=\"cosine, per step\",\n", " batch_size=BATCH, epochs=EPOCHS, eval_epochs=list(EVAL_EPOCHS),\n", " params=sum(p.numel() for p in model.parameters()), seed=seed)\n", " if isinstance(model, AttentionReader):\n", " config.update(attention=\"scaled dot-product\", d_k=64, d_v=128)\n", " run = wandb.init(project=\"cv-lab8\", name=name, config=config)\n", " history = []\n", " for epoch in range(1, EPOCHS + 1):\n", " strips = StripDataset(train_imgs, train_labels, N_STRIPS, lengths,\n", " seed=1000 + epoch)\n", " loader = DataLoader(strips, batch_size=BATCH, shuffle=True, collate_fn=collate,\n", " generator=g, drop_last=True)\n", " model.train()\n", " total = 0.0\n", " for images, dec_in, tgt in loader:\n", " images, dec_in, tgt = images.to(DEV), dec_in.to(DEV), tgt.to(DEV)\n", " loss = masked_loss(model(images, dec_in), tgt)\n", " opt.zero_grad()\n", " loss.backward()\n", " opt.step()\n", " sched.step()\n", " total += loss.item()\n", " row = {\"epoch\": epoch, \"train/loss\": total / len(loader)}\n", " msg = f\"epoch {epoch:2d} train loss {row['train/loss']:.4f}\"\n", " if epoch in EVAL_EPOCHS:\n", " preds, targets, _ = predict(model, test_set)\n", " row.update({\"test/exact\": exact_string_accuracy(preds, targets),\n", " \"test/cer\": cer(preds, targets)})\n", " msg += f\" test exact {row['test/exact']:.4f} CER {row['test/cer']:.4f}\"\n", " wandb.log(row, step=epoch)\n", " history.append(row)\n", " print(msg)\n", " run.finish()\n", " return model, history" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Q4.1.** Before you train them: both readers train on the same strips of one to eight digits, with the same encoder and the same recipe. Which reader do you expect to lose more exact-string accuracy as the strips lengthen from one digit to eight, and why? Commit to an answer here; Q4.3 checks it.\n", "\n", "*Answer:*\n", "\n", "**Q4.2.** Train both readers with seed 0, the fixed-context reader first: `train_reader(FixedReader, 0, ...)` and `train_reader(AttentionReader, 0, ...)`, each with a meaningful run name. Plot the training loss against epoch and the test exact-string accuracy against the epochs at which it was measured, with both readers on each plot." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Q4.3.** Evaluate both readers on `test_set`. Report, for each reader, its parameter count and the exact-string accuracy and CER over lengths one to five together and over six to eight together, as a table, and plot exact-string accuracy and CER against length for both readers. Was your prediction in Q4.1 right? Describe where the two readers differ and where they do not. The attention reader has more parameters, so say what this comparison can attribute to attention and what it cannot." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "## 5. Inspecting the alignments\n", "\n", "`show_attention` reads one strip with the attention reader and draws two panels that share a horizontal axis: the strip with each digit's box outlined and labelled, and beneath it the attention weights, one row per output step and one cell per column of $F$, 16 pixels wide. All the rows of one strip share a colour scale, from zero to the largest weight in any of its rows, so a pale row is a row whose weight is spread out, not a row that has been rescaled. The row labels are the symbols the reader wrote." ] }, { "cell_type": "code", "metadata": {}, "source": [ "def show_attention(model, dataset, i, title=\"\"):\n", " \"\"\"Read strip i with an attention reader; draw the strip and its weights.\"\"\"\n", " img, _ = dataset[i]\n", " pred, alpha = generate(model, img.unsqueeze(0).to(DEV))\n", " pred, alpha = pred[0].tolist(), alpha[0].cpu().numpy()\n", " steps = next((k + 1 for k, y in enumerate(pred) if y == EOS), len(pred))\n", " pred, alpha = pred[:steps], alpha[:steps]\n", " width = img.shape[-1]\n", " fig, (top, bottom) = plt.subplots(\n", " 2, 1, figsize=(9, 1.6 + 0.3 * steps), sharex=True, layout=\"constrained\",\n", " gridspec_kw=dict(height_ratios=[1.3, 0.3 * steps]))\n", " top.imshow(img[0], cmap=\"gray\", extent=(0, width, H, 0))\n", " for d, x in zip(dataset.digits[i], dataset.xs[i]):\n", " top.add_patch(plt.Rectangle((x, 0.5), 28, H - 1, fill=False, ec=\"tab:orange\", lw=1))\n", " top.text(x + 14, -1, str(d), ha=\"center\", va=\"bottom\", fontsize=8, color=\"tab:orange\")\n", " top.set_yticks([])\n", " top.set_title(f\"{title}target {''.join(map(str, dataset.digits[i]))} \"\n", " f\"read {''.join(map(str, clean(pred)))}\", loc=\"left\", fontsize=9, pad=12)\n", " im = bottom.imshow(alpha, cmap=\"viridis\", vmin=0, vmax=alpha.max(), aspect=\"auto\",\n", " interpolation=\"nearest\", extent=(0, width, steps - 0.5, -0.5))\n", " for x in dataset.xs[i]:\n", " bottom.axvline(x, color=\"white\", lw=0.6, alpha=0.6)\n", " bottom.set_yticks(range(steps), [SYMBOLS[y] for y in pred], fontsize=8)\n", " bottom.set_ylabel(\"written\")\n", " bottom.set_xlabel(\"strip position (pixels); one attention cell per 16 pixels\")\n", " fig.colorbar(im, ax=[top, bottom], fraction=0.025, pad=0.01, label=\"weight\")\n", " plt.show()" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Q5.1.** Find a strip of seven or eight digits in `test_set` that the attention reader reads exactly, and a strip of six to eight digits that it reads wrongly, and display both with `show_attention`. For the correct reading, describe where the weight falls at each output step, including the step that writes ``, and how large the largest weight in each row is. For the failure, find the first wrong symbol and describe the weights at that step and the steps after it. Say what the maps show and what they leave open: which column state received the largest weight in a context, compared with which part of the image the written symbol depended on." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "**Q5.2.** Measure the alignment over the whole of `test_set`. `generate` keeps decoding a strip until every strip in its batch has finished, so a prediction can contain symbols after its own ``; discard them, since the reader has stopped. For each strip, list the digits written before the first ``, together with the generation step at which each was written (skip any `` or ``, as `clean` does). Then, for each target position $k$ (counting from 0):\n", "\n", "- if the reader wrote fewer than $k + 1$ digits, count position $k$ as **missing** and do not score it;\n", "- otherwise take the weights at the step that wrote its $k$-th digit, and the column $c$ (counting from 0) where they are largest. Column $c$ corresponds to strip positions $16c$ to $16c + 15$, so take its centre to be pixel $16c + 8$, and count the position as **aligned** if that centre lies inside the box of the target's $k$-th digit, from `xs[k]` to `xs[k] + 27`.\n", "\n", "Report the number of missing positions, the share of aligned positions among the digits read correctly (the $k$-th digit written equals the target's $k$-th digit) and among those read wrongly, and the median of the largest weight over the correctly read digits. What does a high share establish about how the reader works, and what does it not?" ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "**Q5.3.** An attention weight is part of the computation. A claim that a prediction depends on part of the image needs an intervention and a measured change in behaviour. Take the first 1,000 strips of six to eight digits in `test_set` that both readers read exactly, so that the two readers are tested on the same strips. In each, choose one position $k$ at random with `np.random.default_rng(5)`, replace that digit's source image by a random MNIST test image of a different class, redraw the strip with `draw_strip` at the same box positions, and give the same new strips to both readers. Report three proportions for each reader.\n", "\n", "1. The proportion of strips read exactly as the new string.\n", "2. The proportion whose reading has a $k$-th digit equal to the new digit, whatever happens elsewhere.\n", "3. The proportion whose reading changed in length or in any digit other than the $k$-th.\n", "\n", "What does this experiment establish that Q5.2 does not, and what does it leave open? This experiment does not require attention weights." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "## 6. Beyond the training lengths\n", "\n", "Neither reader has seen a strip of more than eight digits. Both can process one: the canvas holds twelve, the encoder produces 21 column states whatever the strip holds, and `generate` allows 14 steps. Being able to accept a longer input is not the same as reading it correctly, and this section measures the difference.\n", "\n", "**Q6.1.** Predict, before you run it, how each reader will do on strips of nine to twelve digits, and whether attention will help there. Then evaluate both readers on `long_set`, and report, for each reader and each length from nine to twelve, the exact-string accuracy, the CER and how often each output length occurs. Classify each strip as Lab 7 did, by comparing cleaned lengths: correct, *missed* digits (shorter than the target), *extra* digits (longer) or *incorrect* digits (the same length but different). Also report the position of the first wrong digit, counting from the first digit as position 0, as a count per position. Which reader does better, by which measure, and why can exact-string accuracy and CER rank the readers differently at these lengths?" ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "**Q6.2.** Display the attention maps of three strips from `long_set` that the attention reader reads wrongly, including at least one of twelve digits, and compute the Q5.2 alignment share on `long_set`. Describe where the weights go at the steps before the first error and at and after it, and compare the length of what the reader wrote with the length of the target. Propose an explanation for the failure, and say what evidence would test it. Attention weights show how the column states were mixed into each context; they do not on their own show why the decoder wrote what it wrote, so state which parts of your explanation the maps support and which they do not." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# YOUR CODE HERE" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*Answer:*\n", "\n", "## 7. Optional explorations\n", "\n", "These are not marked, and neither is needed to complete the experiment above.\n", "\n", "- The caption-attention widget in Chapter 8 shows recorded attention for a model generating words about a photograph. Compare its maps with this reader's: how concentrated is the weight in each, and what does a caption word have that corresponds to a digit's box, if anything?\n", "- Use the chapter's position-code explorer to test the same-offset relationship, or implement sinusoidal position codes from the chapter's formula and plot the dot product between the codes of positions $p$ and $q$ against $p - q$.\n", "\n", "## 8. W&B runs\n", "\n", "Paste links to your W&B runs below, three at minimum, each with a meaningful name and a config that records the following.\n", "\n", "- **The one-batch ritual:** reader, state and embedding widths, $d_k$ and $d_v$, number of strips in the batch, number of steps, optimiser, learning rate, parameter count and seed.\n", "- **The two readers:** the configs `train_reader` writes, which carry the reader, the widths, the training lengths, strips per epoch, optimiser, learning rate and schedule, batch size, epochs, parameter count and seed. Each logs the training loss per epoch, and test exact-string accuracy and CER at the epochs in `EVAL_EPOCHS`.\n", "\n", "*W&B run links:*\n", "\n", "## 9. Before you submit\n", "\n", "- [ ] **Runtime ▸ Restart session and run all** on a GPU runtime, then read every output. The full re-run downloads MNIST, runs the ritual and trains two readers. A reference run took about 38 minutes on a Colab T4; budget fifty, and longer on a slower GPU.\n", "- [ ] Fill in the group members; check that every member can explain every cell.\n", "- [ ] Complete every *Answer:* cell, and paste the W&B links into Section 8.\n", "- [ ] **File ▸ Download ▸ Download .ipynb**; one member submits it on Moodle by the deadline shown there." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 4 }