{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Lab 7 — Reading a sequence from an image\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-07.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 7 of the course book](https://courses.ms.wits.ac.za/~richard/cv/book/chapters/07-recurrent-networks.html). In every network you have trained so far, the number of outputs was set by the architecture or by the size of the image: one label per image, one score per anchor, one value per pixel. This one writes a sequence and chooses its length itself. You build the image-to-sequence model from the end of the chapter, a CNN whose feature map is projected into the first state of a GRU decoder, and train it to read strips of one to five handwritten digits. After twelve epochs the reference solution reads about 97% of the test strips exactly, with a character error rate of about 0.01; three reference runs, with seeds 0, 1 and 2, gave exact-string accuracies of 0.975, 0.971 and 0.975. Your run will differ in detail: the error counts, the commonest confusions and the behaviour on longer strings vary from run to run more than the headline accuracy does. Then you show it strips of six to eight digits, which it has never seen, and work out why it fails on them.\n",
        "\n",
        "This lab's rung of the skills ladder is **variable-length batches: padding and a masked loss**. Strings of different lengths cannot be stacked into one tensor until the short ones are padded to the length of the longest, and once they are padded, the padding positions must be kept out of the loss. Otherwise part of every batch's loss comes from predicting padding, which is not a task the model needs to learn. You write the batching function that pads, and a loss that ignores the padding but still treats `<eos>` as a genuine target, and you check both against a calculation you do by hand.\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-07.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-lab7`."
      ]
    },
    {
      "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: 60,000 training and 10,000 test images of single handwritten digits, each $28 \\times 28$ pixels, white on black (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 has thirteen symbols. The digits 0 to 9 are their own indices; index 10 is `<sos>`, the start symbol the decoder reads first; index 11 is `<eos>`, the end symbol the decoder writes to stop; and index 12 is `<pad>`, which fills the unused positions of short sequences in a batch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "SOS, EOS, PAD = 10, 11, 12\n",
        "VOCAB = 13\n",
        "SYMBOLS = [str(d) for d in range(10)] + [\"<sos>\", \"<eos>\", \"<pad>\"]\n",
        "\n",
        "\n",
        "def show_tokens(seq):\n",
        "    return \" \".join(SYMBOLS[int(t)] for t in seq)\n",
        "\n",
        "\n",
        "print(show_tokens([SOS, 4, 2, 7, EOS, PAD]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The scoring functions below are supplied, so that your time goes on interpreting the errors rather than on implementing edit distance. Section 4 explains them; for now, `clean` turns any token sequence into a list of digits by stopping at the first `<eos>` and dropping `<sos>` and `<pad>`, and `exact_string_accuracy` compares cleaned predictions with cleaned targets. `edit_ops` lists the individual edits behind an edit distance, for Q4.4."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "def clean(seq):\n",
        "    \"\"\"Digits only: stop at the first <eos>, drop <sos> and <pad>.\"\"\"\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 edit_ops(pred, target):\n",
        "    \"\"\"One minimal alignment of the cleaned prediction against the cleaned target.\n",
        "\n",
        "    Returns a list of operations, one per edit, in target order:\n",
        "      (\"sub\",   true digit, predicted digit, index in prediction)\n",
        "      (\"miss\",  true digit, None,            index in prediction where it is missing)\n",
        "      (\"extra\", None,       predicted digit, index in prediction)\n",
        "    len(edit_ops(p, t)) == edit_distance(clean(p), clean(t)).\n",
        "    \"\"\"\n",
        "    p, t = clean(pred), clean(target)\n",
        "    n, m = len(t), len(p)\n",
        "    D = [[0] * (m + 1) for _ in range(n + 1)]\n",
        "    for i in range(n + 1):\n",
        "        D[i][0] = i\n",
        "    for j in range(m + 1):\n",
        "        D[0][j] = j\n",
        "    for i in range(1, n + 1):\n",
        "        for j in range(1, m + 1):\n",
        "            D[i][j] = min(D[i - 1][j] + 1, D[i][j - 1] + 1,\n",
        "                          D[i - 1][j - 1] + (t[i - 1] != p[j - 1]))\n",
        "    out, i, j = [], n, m\n",
        "    while i > 0 or j > 0:\n",
        "        if i > 0 and j > 0 and D[i][j] == D[i - 1][j - 1] + (t[i - 1] != p[j - 1]):\n",
        "            if t[i - 1] != p[j - 1]:\n",
        "                out.append((\"sub\", t[i - 1], p[j - 1], j - 1))\n",
        "            i, j = i - 1, j - 1\n",
        "        elif i > 0 and D[i][j] == D[i - 1][j] + 1:\n",
        "            out.append((\"miss\", t[i - 1], None, j))\n",
        "            i -= 1\n",
        "        else:\n",
        "            out.append((\"extra\", None, p[j - 1], j - 1))\n",
        "            j -= 1\n",
        "    return out[::-1]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Warm-up: copying a digit string\n",
        "\n",
        "Before the model reads an image, it reads a string. The task is to copy five random digits: a GRU encoder reads `4 2 7 5 9`, and a GRU decoder writes `4 2 7 5 9 <eos>`. This is the encoder–decoder of the chapter's translation example with the same vocabulary on both sides. Every string here has exactly five digits, so every batch is already rectangular and no padding is needed yet; Section 2 removes that restriction.\n",
        "\n",
        "The model follows the chapter's notation. The encoder starts from a zero state and reads the source embeddings; its final state $h_T$ is projected to the decoder's first state, $s_0 = \\tanh(W_{\\text{init}} h_T + b_{\\text{init}})$. The decoder then computes $s_t = \\operatorname{GRU}(E[y_{t-1}], s_{t-1})$ and the logits $\\ell_t = W_{\\text{out}} s_t + b_{\\text{out}}$. With teacher forcing, $y_{t-1}$ is the correct previous symbol, so the decoder's inputs are `<sos> 4 2 7 5 9` and its targets are `4 2 7 5 9 <eos>`. `copy_batch` builds all three."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "def copy_batch(n, length=5, generator=None):\n",
        "    \"\"\"Random digit strings with the decoder input and target for each.\"\"\"\n",
        "    src = torch.randint(0, 10, (n, length), generator=generator)\n",
        "    sos = torch.full((n, 1), SOS, dtype=torch.long)\n",
        "    eos = torch.full((n, 1), EOS, dtype=torch.long)\n",
        "    return src, torch.cat([sos, src], 1), torch.cat([src, eos], 1)\n",
        "\n",
        "\n",
        "src, dec_in, tgt = copy_batch(1)\n",
        "print(\"source\", show_tokens(src[0]), \"| decoder input\", show_tokens(dec_in[0]),\n",
        "      \"| target\", show_tokens(tgt[0]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q1.1.** Implement `Copier(hidden=128, emb=32)`. It needs a source embedding table and an `nn.GRU` encoder, the projection $W_{\\text{init}}$ as an `nn.Linear(hidden, hidden)`, a separate target embedding table $E$ and `nn.GRU` decoder, and the output layer $W_{\\text{out}}$ as an `nn.Linear(hidden, VOCAB)`. Use `batch_first=True` for both GRUs. Give it two methods: `encode(src)`, which returns $s_0$ in the shape `nn.GRU` expects for an initial state, $(1, B, \\text{hidden})$; and `forward(src, dec_in)`, which runs the decoder over all of `dec_in` at once from $s_0$ and returns logits of shape $(B, T, 13)$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "class Copier(nn.Module):\n",
        "    # YOUR CODE HERE\n",
        "    ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q1.2.** Write `greedy_copy(model, src, max_len=12)`, which generates without the target. Start every row from `<sos>` and $s_0$, run the decoder one step at a time, take the highest-scoring symbol at each step, and feed it in as the next input. Run all `max_len` steps and return the chosen symbols as a tensor of shape $(B, \\text{max\\_len})$; `clean` discards everything after the first `<eos>`. Put the model in evaluation mode while generating and back in training mode afterwards, and do not track gradients."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "@torch.no_grad()\n",
        "def greedy_copy(model, src, max_len=12):\n",
        "    # YOUR CODE HERE\n",
        "    ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q1.3.** Before you train it: what cross-entropy do you expect at the first step from an untrained `Copier`, and why? Then call `seed_everything(0)`, build a `Copier`, and train it for 1,500 steps with Adam at $10^{-3}$, each step on a fresh `copy_batch(256)`, using `F.cross_entropy` between the logits and the targets under teacher forcing. Log the loss to W&B in a run named for the copying task, with a config that records the widths, optimiser, learning rate, batch size, number of steps and seed. Every 250 steps, and at step 1, print the loss and the free-running exact-string accuracy on a fixed evaluation set of 1,000 strings, `copy_batch(1000, generator=torch.Generator().manual_seed(99))`. Finally, print three generated sequences in full with `show_tokens`, before cleaning."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q1.4.** The copier has only ever seen strings of five digits. Predict what it writes when it is given a string of three digits, and a string of seven. Then test both on 1,000 strings each and report the exact-string accuracy and how often each output length occurs (`np.bincount` of `len(clean(p))` is enough), with one example of each. What has the decoder learned about when to write `<eos>`?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "## 2. Digit strips, padding and the masked loss\n",
        "\n",
        "The reader in Section 3 sees an image instead of a source string. The supplied generator draws $k$ digits at random from a set of MNIST images and places them left to right on a blank canvas 28 pixels high and 224 pixels wide, which is wide enough for eight digits. Neighbouring digit boxes overlap by between 0 and 6 pixels, which MNIST's blank margins absorb, and the whole string starts at a random horizontal position. Every canvas has the same size, so a batch of images stacks without further work: the image is padded with blank canvas when it is generated.\n",
        "\n",
        "`StripDataset` builds `n` strips with equally many of each length in `lengths`, fixed by `seed`. Item $i$ is a pair: the strip as a float tensor of shape $(1, 28, 224)$ in $[0, 1]$, and its digits as a `long` tensor with no special symbols."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "H, W = 28, 224\n",
        "GAP_RANGE = (-6, 0)\n",
        "\n",
        "\n",
        "def make_strip(images, labels, length, rng):\n",
        "    \"\"\"One strip of `length` digits on a blank 28 x 224 canvas.\n",
        "\n",
        "    Digits are drawn at random from `images` (uint8, N x 28 x 28) and placed\n",
        "    left to right. Neighbouring 28-pixel boxes overlap by 0 to 6 pixels, which\n",
        "    MNIST's blank margins absorb, and the whole string starts at a random\n",
        "    horizontal offset. Returns the canvas (uint8, 28 x 224) and the digit list.\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",
        "    width = 28 * length + int(gaps.sum())\n",
        "    x = int(rng.integers(0, W - width + 1))\n",
        "    canvas = np.zeros((H, W), dtype=np.uint8)\n",
        "    for k, i in enumerate(idx):\n",
        "        np.maximum(canvas[:, x:x + 28], images[i], out=canvas[:, x:x + 28])\n",
        "        if k < length - 1:\n",
        "            x += 28 + int(gaps[k])\n",
        "    return canvas, [int(labels[i]) for i in idx]\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, 224) in [0, 1]\n",
        "    and a long tensor holding the digits only, with no special symbols.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, images, labels, n, lengths, seed):\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, W), dtype=np.uint8)\n",
        "        self.digits = []\n",
        "        for j, L in enumerate(lens):\n",
        "            self.images[j], d = make_strip(images, labels, int(L), rng)\n",
        "            self.digits.append(d)\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",
        "train_set = StripDataset(train_imgs, train_labels, 60_000, range(1, 6), seed=1000)\n",
        "test_set = StripDataset(test_imgs, test_labels, 5_000, range(1, 6), seed=1234)\n",
        "long_set = StripDataset(test_imgs, test_labels, 3_000, range(6, 9), seed=4321)\n",
        "print(len(train_set), 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 five and `long_set` 1,000 of each length from six to eight, all from MNIST test digits and identical on every run. The model never trains on `long_set`; Section 5 uses it.\n",
        "\n",
        "**Q2.1.** Display eight strips from `train_set` and four from `long_set`, each titled with its digits. Then, before you compute it, predict the share of training strips that have ink in pixel column 14, near the left edge; in column 112, the middle; and in column 210, near the right edge. Compute the three shares over `train_set`. Why does the generator start each string at a random position rather than at the left edge, given what Section 5 will test?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "A batch of strips has strings of different lengths, so the decoder inputs and targets have to be padded to the longest string in the batch before they can be stacked. For the strings `4 2 7` and `5 1`, the padded batch is\n",
        "\n",
        "| | step 1 | step 2 | step 3 | step 4 |\n",
        "|---|---|---|---|---|\n",
        "| decoder input, string 1 | `<sos>` | `4` | `2` | `7` |\n",
        "| target, string 1 | `4` | `2` | `7` | `<eos>` |\n",
        "| decoder input, string 2 | `<sos>` | `5` | `1` | `<pad>` |\n",
        "| target, string 2 | `5` | `1` | `<eos>` | `<pad>` |\n",
        "\n",
        "The `<eos>` of a short string is a genuine target, because the model has to learn to stop there. The `<pad>` after it is not, and neither is the `<pad>` in the decoder input, whose prediction is scored against a `<pad>` target and is therefore excluded too.\n",
        "\n",
        "**Q2.2.** Write `collate(batch)` for a `DataLoader`. It receives a list of `(image, digits)` pairs and returns three tensors: the images stacked to $(B, 1, 28, 224)$; the decoder inputs, each `<sos>` followed by the digits; and the targets, each the digits followed by `<eos>`. Pad the decoder inputs and the targets with `PAD` to the same length $T$, one more than the longest string in the batch (`pad_sequence` with `batch_first=True` and `padding_value=PAD` does this). Print a batch of four with `show_tokens`, and check it against the table above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "def collate(batch):\n",
        "    # YOUR CODE HERE\n",
        "    ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q2.3.** Write `masked_loss(logits, tgt)`, the mean cross-entropy over every target position that is not `<pad>`, using the `ignore_index` argument of `F.cross_entropy`. Then do three things.\n",
        "\n",
        "1. Before you count: training strips have one to five digits in equal numbers, and batches hold 128 strips. Predict roughly what share of the target positions in a batch are `<pad>`. Then measure it over one pass through `train_set`.\n",
        "2. Check `masked_loss` on eight strips with random logits (`torch.randn` of shape $(8, T, 13)$, after `torch.manual_seed(0)`) against a loop that computes the cross-entropy of each strip's real positions separately, with `reduction=\"none\"`, and averages them all. Print both, the mean of the eight per-strip means, and the unmasked `F.cross_entropy` on the same logits.\n",
        "3. Answer: which of the two averages does `ignore_index` compute, and how do they weight long and short strings differently? What would the decoder be trained to do at the padded positions if you left `ignore_index` out, and why does that also make the reported loss harder to interpret?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "def masked_loss(logits, tgt):\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",
        "## 3. The reader: a CNN encoder and a GRU decoder\n",
        "\n",
        "This is the model of the chapter's \"Image captioning\" section, with a strip in place of a photograph. A CNN of four blocks, each a $3 \\times 3$ convolution with `padding=1` and `bias=False`, batch norm, ReLU and a $2 \\times 2$ max-pool, at 32, 64, 128 and 128 channels, turns the strip into a feature map $F$. Here $F$ has 128 channels, a height of 1 and a width of 14. `F.flatten(1)` lists the 14 positions of the first channel from left to right, then those of the second channel, and so on, so features from the left of the strip and from the right occupy different entries of $\\operatorname{vec}(F)$. Listing the channels at each position in turn would also keep the positions apart; what matters is that the order is fixed. A linear layer and tanh project that vector to the decoder's first state,\n",
        "\n",
        "$$\n",
        "s_0=\\tanh\\!\\left(W_{\\mathrm{init}}\\operatorname{vec}(F)+b_{\\mathrm{init}}\\right),\n",
        "$$\n",
        "\n",
        "and from there the decoder is the same as the copier'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 and a state width of 256.\n",
        "\n",
        "**Q3.1.** Implement `Reader(hidden=256, emb=32)` with the same two methods as the copier, `encode(images)` returning $s_0$ as $(1, B, 256)$ and `forward(images, dec_in)` returning logits of shape $(B, T, 13)$. Give both a `trace=False` argument that, when true, prints the shapes of $F$, $\\operatorname{vec}(F)$, $s_0$, the decoder states and the logits. You need the shape of $F$ to size $W_{\\text{init}}$, so work it out from the pooling before you write that line."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "class Reader(nn.Module):\n",
        "    # YOUR CODE HERE\n",
        "    ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q3.2.** Before you run it: write down the shape of $F$ for a batch of four strips, and compute the number of parameters in $W_{\\text{init}}$ and $b_{\\text{init}}$ by hand. Then run the trace on a batch of four from `collate`, print the parameter count of each part (CNN, projection, embedding, GRU, output layer) and the total, and check both. Which part holds the most parameters, and what does that say about where the image's information has to pass?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q3.3.** The ritual, from Lab 2 and expected before every real run since: take the first 32 strips of `train_set` as one batch, call `seed_everything(0)`, and drive a fresh `Reader` to memorise them with Adam at $10^{-3}$ for 300 steps, using teacher forcing and `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": [
        "**Q3.4.** Write `generate(model, images, max_len=12)`, the reader's version of `greedy_copy`: start from `<sos>` and $s_0$, feed back the highest-scoring symbol at each step, and return the symbols chosen. Stop early once every row in the batch has produced an `<eos>`, so the returned tensor has at most `max_len` columns. Then write `predict(model, dataset)`, which runs `generate` over a whole dataset in batches of 500 and returns two lists of token lists, the predictions and the targets, ready for the scoring functions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "@torch.no_grad()\n",
        "def generate(model, images, max_len=12):\n",
        "    # YOUR CODE HERE\n",
        "    ...\n",
        "\n",
        "\n",
        "def predict(model, dataset, batch_size=500):\n",
        "    # YOUR CODE HERE\n",
        "    ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q3.5.** Train the reader. Write `train_reader(lengths, seed, name)`, because Section 5 runs it again, and use this recipe.\n",
        "\n",
        "- `seed_everything(seed)`, then a fresh `Reader()` and Adam at $10^{-3}$, with `torch.optim.lr_scheduler.CosineAnnealingLR(opt, 12 * 468)` stepped after every batch.\n",
        "- Twelve epochs. At the start of epoch $e$, build a fresh `StripDataset(train_imgs, train_labels, 60_000, lengths, seed=1000 + e)`, so that every epoch sees new strips drawn from the training digits, and a loader with `batch_size=128, shuffle=True, drop_last=True`, `collate_fn=collate` and `generator=torch.Generator().manual_seed(seed)` created once before the first epoch. That gives 468 batches per epoch.\n",
        "- Teacher forcing and `masked_loss`.\n",
        "- After every epoch, the free-running exact-string accuracy and CER on `test_set`, logged to W&B with the mean training loss.\n",
        "\n",
        "The config must carry the state and embedding widths, the training lengths, the number of strips per epoch, the optimiser, learning rate and schedule, the batch size, the epoch count, the parameter count and the seed. This lab has no separate validation split, so the per-epoch test numbers monitor the test split for illustration only: keep the prescribed settings fixed and do not use the curves to choose them. Call `train_reader(range(1, 6), 0, ...)` and plot the training loss and the test exact-string accuracy against epoch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Free-running evaluation\n",
        "\n",
        "From here on, results are free-running unless a question says otherwise: the model sees the image and its own previous outputs, never the target. Q4.3 compares this with teacher forcing. Evaluate in evaluation mode (`model.eval()`) and without gradients (`torch.no_grad()`), so that batch norm uses its stored statistics and does not update them; `generate` already does both. Two measures score it. **Exact-string accuracy** is the proportion of images for which the whole digit string is correct, and it is the main result. **Character error rate** (CER) is the minimum number of digit substitutions, deletions and insertions that turn the prediction into the target, its *edit distance*, divided by the number of target digits. Over a test set, the edits are summed over all strips and divided by the total number of target digits, so a long string counts for more than a short one. Lower CER is better, and it can exceed 1 when a prediction contains many extra digits. Both measures score cleaned strings: generation stops at the first `<eos>` or after `max_len` steps, and `<sos>`, `<eos>` and `<pad>` are removed before comparison.\n",
        "\n",
        "**Q4.1.** For the target `427`, compute by hand the edit distance and CER of the predictions `47`, `417` and `4227`, and say which operation each one needs. Then check your answers with `edit_distance`, `cer` and `exact_string_accuracy`. Find a prediction for the target `4` whose CER is above 1. Finally, show what `clean` returns for `[4, 2, EOS, 7]` and for `[SOS, 4, PAD]`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q4.2.** Report exact-string accuracy and CER on `test_set` overall and for each string length from one to five, as a table and as a plot of both measures against length. Describe how the two curves move with length in your run; they need not be monotonic, and a dip or a crossing is a result to report, not to smooth over. Then compare your exact-string accuracies with a simple model. If every target digit were wrong independently with probability $p$, and every error were a substitution, a string of length $L$ would be exactly right with probability $(1 - p)^L$. CER is an edit rate rather than a probability, and it estimates $p$ only when substitutions account for nearly all the edits, which Q4.4 lets you check. Using each length's CER as $p$, compute $(1 - \\text{CER})^L$ beside the measured value, and say where the two differ and what that suggests about how the errors are distributed across strips."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q4.3.** Teacher-forced accuracy is what you get when every step is given the true previous digit, as in training. Measure it in evaluation mode and under `torch.no_grad()`: run `model(images, dec_in)` on the padded batches, take the argmax at every position, and count a strip as correct when every non-padding target position, `<eos>` included, matches. That criterion is stricter than `exact_string_accuracy`, which cleans the prediction first and so forgives a stray `<sos>` or `<pad>` and a missing `<eos>`. For a like-for-like comparison, also compute a **strict free-running** accuracy: a generated sequence counts only if its first tokens are exactly the target digits followed at once by `<eos>`. Predict whether teacher-forced and strict free-running accuracy will differ, then report all three numbers and the number of strips on which the strict and cleaned criteria disagree. Explain when teacher-forced and strict free-running accuracy must agree, and what could make them differ. Would per-digit accuracy behave the same way?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q4.4.** Write `error_kind(pred, target)`, which sorts a prediction by comparing cleaned lengths: **correct**; **missed** digits (the prediction is shorter than the target); **extra** digits (longer); or **incorrect** digits (the same length, but different). Count each kind on `test_set`, and display up to four examples of each kind that occurs, as images titled with the target and the prediction. Then use the supplied `edit_ops` over all the predictions to count substitutions, missed digits and extra digits, to list the commonest substitutions as (true, predicted) pairs, and to count how many extra digits equal a neighbouring digit in the prediction. Finally, count the incorrect predictions that hold the target's digits in a different order, such as `94` read as `49`, and display a few. Which kind of error dominates? Does it look more like a failure to recognise digits or a failure to keep track of position and count? Treat a reordering as a hypothesis to check against the images: `94` read as `49` could also be two separate recognition errors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "## 5. Longer strings than it was trained on\n",
        "\n",
        "The canvas is 224 pixels wide, so strips of six, seven and eight digits fit on it without any change to the model: the same CNN, the same $W_{\\text{init}}$, the same $s_0$ of 256 numbers. Nothing stops the decoder from writing more than five digits either, because `generate` allows twelve steps. The reader has never been trained on a strip of more than five digits.\n",
        "\n",
        "**Q5.1.** Predict, before you run it, what the reader will write for a strip of eight digits. Then evaluate it on `long_set` and report, for each length from six to eight, the exact-string accuracy, the CER, how often each output length occurs, and the share of strips whose first five digits are read correctly. Compare that last number with the exact-string accuracy at length five in Q4.2. Classify the errors with `error_kind` and report the counts for each length; which kind dominates now, and does the answer depend on the length? Display six strips from `long_set` with their targets and predictions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q5.2.** Q5.1 changed two things at once: the strings are longer than any the reader was trained on, and the same 256 numbers of $s_0$ now have to hold more digits. A control helps to tell these apart. Call `train_reader(range(1, 9), 0, ...)`, the same recipe with training strips of one to eight digits and nothing else in the code changed. Print, for each length from one to eight, the exact-string accuracy and CER, using `test_set` for lengths one to five and `long_set` for six to eight, and for lengths six to eight the output-length counts and the `error_kind` counts. No plots are required."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q5.3.** What does the control in Q5.2 rule out as an explanation of the failure in Q5.1, and what does it leave unresolved? Training on lengths one to eight changed more than the longest length: the mixture of lengths, the number of supervised positions per strip and how much of the canvas is inked all changed too, so say what the comparison can and cannot attribute. Then consider what neither experiment varied: $s_0$ has 256 numbers whether the strip holds one digit or eight. What would you expect to happen as the strips grow longer, or as the state narrows, and what would you measure to find out? Changing the state width also changes the size of $W_{\\text{init}}$ and the GRU, so say what such an experiment could and could not isolate. Chapter 8 opens with this question, under the name of the *fixed-width context bottleneck*.\n",
        "\n",
        "*Answer:*\n",
        "\n",
        "## 6. W&B runs\n",
        "\n",
        "Paste links to your W&B runs below, four at minimum, each with a meaningful name and a config that records the following.\n",
        "\n",
        "- **The copying warm-up:** state and embedding widths, string length, optimiser, learning rate, batch size, number of steps, parameter count and seed.\n",
        "- **The one-batch ritual:** state and embedding widths, number of strips in the batch, number of steps, optimiser, learning rate, parameter count and seed.\n",
        "- **The two readers, trained on lengths one to five and one to eight:** state and embedding widths, training lengths, strips per epoch, optimiser, learning rate and schedule, batch size, epochs, parameter count and seed. Each logs the training loss, test exact-string accuracy and test CER per epoch.\n",
        "\n",
        "*W&B run links:*\n",
        "\n",
        "## 7. 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, trains the copier, runs the ritual and trains two readers. A reference run took about 16 minutes on a Colab T4; budget twenty, 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 6.\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
}