{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-1 fw=all -->\n",
    "\n",
    "# Utility Functions and Classes\n",
    "\n",
    "\n",
    "This section contains the implementations of utility functions and classes used in this book."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-1"
   },
   "source": [
    "import inspect\n",
    "import collections\n",
    "from d2l import torch as d2l\n",
    "from IPython import display\n",
    "from torch import nn"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-2 fw=all -->\n",
    "\n",
    "Hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-2"
   },
   "source": [
    "@d2l.add_to_class(d2l.HyperParameters)\n",
    "def save_hyperparameters(self, ignore=None):\n",
    "    \"\"\"Save function arguments into class attributes.\"\"\"\n",
    "    ignore = [] if ignore is None else ignore\n",
    "    frame = inspect.currentframe().f_back\n",
    "    _, _, _, local_vars = inspect.getargvalues(frame)\n",
    "    self.hparams = {k:v for k, v in local_vars.items()\n",
    "                    if k not in set(ignore+['self']) and not k.startswith('_')}\n",
    "    for k, v in self.hparams.items():\n",
    "        setattr(self, k, v)"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-3 fw=all -->\n",
    "\n",
    "Progress bar."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-3"
   },
   "source": [
    "\n",
    "import io, os, queue, threading, time\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def draw(self, x, y, label, every_n=1):\n",
    "    \"\"\"Schedule the point (x, y) to be plotted under `label`.\n",
    "\n",
    "    This call is *asynchronous*: it appends the point to an internal queue\n",
    "    and returns immediately, so the (possibly compiled) training loop never\n",
    "    blocks on a device-to-host transfer or on matplotlib. `y` may be a number\n",
    "    or a zero-argument callable returning one; if it is a callable, the\n",
    "    conversion runs on the background drawing thread, not on the caller's.\n",
    "    `every_n` averages each group of `n` raw points into one plotted point.\"\"\"\n",
    "    self._start_drawer()\n",
    "    try:\n",
    "        self._queue.put_nowait((x, y, label, every_n))\n",
    "    except queue.Full:\n",
    "        pass  # drop the point rather than stall the training loop\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _start_drawer(self):\n",
    "    if getattr(self, '_drawer', None) is not None:\n",
    "        return\n",
    "    self._queue = queue.Queue(maxsize=1000)\n",
    "    self._lock = threading.Lock()\n",
    "    self._handle = None\n",
    "    self._last = 0.0\n",
    "    self._running = True\n",
    "    # Emit live frames only outside the headless book build, so an executed\n",
    "    # notebook records exactly one (final) figure in the committed store.\n",
    "    self._live = self.display and os.environ.get('D2L_NB_CAPTURE') != '1'\n",
    "    self._drawer = threading.Thread(target=self._drain, daemon=True)\n",
    "    self._drawer.start()\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _drain(self):\n",
    "    while self._running or not self._queue.empty():\n",
    "        try:\n",
    "            batch = [self._queue.get(timeout=0.1)]\n",
    "        except queue.Empty:\n",
    "            continue\n",
    "        while True:  # coalesce everything currently queued\n",
    "            try: batch.append(self._queue.get_nowait())\n",
    "            except queue.Empty: break\n",
    "        for x, y, label, every_n in batch:\n",
    "            self._accumulate(x, y() if callable(y) else y, label, every_n)\n",
    "        if self._live and time.time() - self._last > 0.1:\n",
    "            self._render(thread=True)\n",
    "            self._last = time.time()\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _accumulate(self, x, y, label, every_n):\n",
    "    Point = collections.namedtuple('Point', ['x', 'y'])\n",
    "    with self._lock:\n",
    "        if not hasattr(self, 'raw_points'):\n",
    "            self.raw_points = collections.OrderedDict()\n",
    "            self.data = collections.OrderedDict()\n",
    "        if label not in self.raw_points:\n",
    "            self.raw_points[label] = []\n",
    "            self.data[label] = []\n",
    "        points = self.raw_points[label]\n",
    "        points.append(Point(x, float(y)))\n",
    "        if len(points) != every_n:\n",
    "            return\n",
    "        mean = lambda v: sum(v) / len(v)\n",
    "        self.data[label].append(Point(mean([p.x for p in points]),\n",
    "                                      mean([p.y for p in points])))\n",
    "        points.clear()\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _plot_lines(self, axes):\n",
    "    with self._lock:\n",
    "        series = [(k, list(v)) for k, v in getattr(self, 'data', {}).items()]\n",
    "    for (k, v), ls, color in zip(series, self.ls, self.colors):\n",
    "        axes.plot([p.x for p in v], [p.y for p in v],\n",
    "                  linestyle=ls, color=color, label=k)\n",
    "    if self.xlim: axes.set_xlim(self.xlim)\n",
    "    if self.ylim: axes.set_ylim(self.ylim)\n",
    "    if not self.xlabel: self.xlabel = self.x\n",
    "    axes.set_xlabel(self.xlabel)\n",
    "    axes.set_ylabel(self.ylabel)\n",
    "    axes.set_xscale(self.xscale)\n",
    "    # A smoothed series exists before it necessarily contains its first point.\n",
    "    # Matplotlib rejects an empty/nonpositive log axis, so keep the temporary\n",
    "    # live frame linear and apply the requested log scale as soon as a positive\n",
    "    # value is available. The final training figure therefore still uses log.\n",
    "    has_positive_y = any(p.y > 0 for _, values in series for p in values)\n",
    "    if self.yscale != 'log' or has_positive_y:\n",
    "        axes.set_yscale(self.yscale)\n",
    "    if series: axes.legend()\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _render(self, thread=False):\n",
    "    if not self.display:\n",
    "        return\n",
    "    if thread:\n",
    "        # Off the main thread: matplotlib's global pyplot state is not\n",
    "        # thread-safe, so render into a private Agg figure and push a PNG.\n",
    "        from matplotlib.figure import Figure\n",
    "        from matplotlib.backends.backend_agg import FigureCanvasAgg\n",
    "        fig = Figure(figsize=self.figsize)\n",
    "        FigureCanvasAgg(fig)\n",
    "        self._plot_lines(fig.add_subplot(111))\n",
    "        buf = io.BytesIO()\n",
    "        fig.savefig(buf, format='png', bbox_inches='tight')\n",
    "        frame = display.Image(data=buf.getvalue())\n",
    "        if self._handle is None:\n",
    "            self._handle = display.display(frame, display_id=True)\n",
    "        else:\n",
    "            self._handle.update(frame)\n",
    "    else:\n",
    "        # Main thread (final frame): render through pyplot so the output is\n",
    "        # captured exactly like the rest of the book's inline figures.\n",
    "        d2l.use_svg_display()\n",
    "        if self.fig is None:\n",
    "            self.fig = d2l.plt.figure(figsize=self.figsize)\n",
    "        self.fig.clf()\n",
    "        self._plot_lines(self.fig.add_subplot(111))\n",
    "        if self._handle is not None:\n",
    "            self._handle.update(self.fig)         # interactive: reuse live slot\n",
    "        else:\n",
    "            # Display, then a trailing clear_output(wait=True): the inline\n",
    "            # backend's automatic end-of-cell render then *replaces* (rather\n",
    "            # than duplicates) this figure, so exactly one image is captured.\n",
    "            display.display(self.fig)\n",
    "            display.clear_output(wait=True)\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def flush(self):\n",
    "    \"\"\"Wait for every scheduled point to be drawn, then render the final\n",
    "    figure on the calling thread. `Trainer.fit` calls this for you; call it\n",
    "    yourself after a bare `draw` loop.\"\"\"\n",
    "    if getattr(self, '_drawer', None) is not None:\n",
    "        self._running = False\n",
    "        self._drawer.join()\n",
    "        self._drawer = None\n",
    "    self._render(thread=False)"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-4 fw=all -->\n",
    "\n",
    "Add FrozenLake environment"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-4"
   },
   "source": [
    "\n",
    "def frozen_lake(seed):\n",
    "    # See https://www.gymlibrary.dev/environments/toy_text/frozen_lake/ to learn more about this env\n",
    "    # How to process env.P.items is adapted from https://sites.google.com/view/deep-rl-bootcamp/labs\n",
    "    import gymnasium as gym\n",
    "\n",
    "    env = gym.make('FrozenLake-v1', is_slippery=False)\n",
    "    env.reset(seed=seed)\n",
    "    env.action_space.seed(seed)\n",
    "    env_info = {}\n",
    "    env_info['desc'] = env.unwrapped.desc  # 2D array specifying what each grid item means\n",
    "    env_info['num_states'] = env.observation_space.n  # Number of observations/states or obs/state dim\n",
    "    env_info['num_actions'] = env.action_space.n  # Number of actions or action dim\n",
    "    # Define indices for (transition probability, nextstate, reward, done) tuple\n",
    "    env_info['trans_prob_idx'] = 0  # Index of transition probability entry\n",
    "    env_info['nextstate_idx'] = 1  # Index of next state entry\n",
    "    env_info['reward_idx'] = 2  # Index of reward entry\n",
    "    env_info['done_idx'] = 3  # Index of done entry\n",
    "    env_info['mdp'] = {}\n",
    "    env_info['env'] = env\n",
    "\n",
    "    for (s, others) in env.unwrapped.P.items():\n",
    "        # others(s) = {a0: [ (p(s'|s,a0), s', reward, done),...], a1:[...], ...}\n",
    "\n",
    "        for (a, pxrds) in others.items():\n",
    "            # pxrds is [(p1,next1,r1,d1),(p2,next2,r2,d2),..].\n",
    "            # e.g. [(0.3, 0, 0, False), (0.3, 0, 0, False), (0.3, 4, 1, False)]\n",
    "            env_info['mdp'][(s,a)] = pxrds\n",
    "\n",
    "    return env_info"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-5 fw=all -->\n",
    "\n",
    "Create environment"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-5"
   },
   "source": [
    "\n",
    "def make_env(name ='', seed=0):\n",
    "    # Input parameters:\n",
    "    # name: specifies a gym environment.\n",
    "    # For Value iteration, only FrozenLake-v1 is supported.\n",
    "    if name == 'FrozenLake-v1':\n",
    "        return frozen_lake(seed)\n",
    "\n",
    "    else:\n",
    "        raise ValueError(f\"{name} env is not supported in this Notebook\")"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-6 fw=all -->\n",
    "\n",
    "Show value function"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-6"
   },
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "def show_value_function_progress(env_desc, V, pi):\n",
    "    # This function visualizes how value and policy changes over time.\n",
    "    # V: [num_iters, num_states]\n",
    "    # pi: [num_iters, num_states]\n",
    "    # How to visualize value function is adapted (but changed) from: https://sites.google.com/view/deep-rl-bootcamp/labs\n",
    "\n",
    "    num_iters = V.shape[0]\n",
    "    fig, ax  = plt.subplots(figsize=(15, 15))\n",
    "\n",
    "    for k in range(V.shape[0]):\n",
    "        plt.subplot(4, 4, k + 1)\n",
    "        plt.imshow(V[k].reshape(4,4), cmap=\"bone\")\n",
    "        ax = plt.gca()\n",
    "        ax.set_xticks(np.arange(0, 5)-.5, minor=True)\n",
    "        ax.set_yticks(np.arange(0, 5)-.5, minor=True)\n",
    "        ax.grid(which=\"minor\", color=\"w\", linestyle='-', linewidth=3)\n",
    "        ax.tick_params(which=\"minor\", bottom=False, left=False)\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])\n",
    "\n",
    "        # LEFT action: 0, DOWN action: 1\n",
    "        # RIGHT action: 2, UP action: 3\n",
    "        action2dxdy = {0:(-.25, 0),1: (0, .25),\n",
    "                       2:(0.25, 0),3: (-.25, 0)}\n",
    "\n",
    "        for y in range(4):\n",
    "            for x in range(4):\n",
    "                action = pi[k].reshape(4,4)[y, x]\n",
    "                dx, dy = action2dxdy[action]\n",
    "\n",
    "                if env_desc[y,x].decode() == 'H':\n",
    "                    ax.text(x, y, str(env_desc[y,x].decode()),\n",
    "                       ha=\"center\", va=\"center\", color=\"y\",\n",
    "                         size=20, fontweight='bold')\n",
    "\n",
    "                elif env_desc[y,x].decode() == 'G':\n",
    "                    ax.text(x, y, str(env_desc[y,x].decode()),\n",
    "                       ha=\"center\", va=\"center\", color=\"w\",\n",
    "                         size=20, fontweight='bold')\n",
    "\n",
    "                else:\n",
    "                    ax.text(x, y, str(env_desc[y,x].decode()),\n",
    "                       ha=\"center\", va=\"center\", color=\"g\",\n",
    "                         size=15, fontweight='bold')\n",
    "\n",
    "                # No arrow for cells with G and H labels\n",
    "                if env_desc[y,x].decode() != 'G' and env_desc[y,x].decode() != 'H':\n",
    "                    ax.arrow(x, y, dx, dy, color='r', head_width=0.2, head_length=0.15)\n",
    "\n",
    "        ax.set_title(\"Step = \"  + str(k + 1), fontsize=20)\n",
    "\n",
    "    fig.tight_layout()\n",
    "    plt.show()"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-7 fw=all -->\n",
    "\n",
    "Show Q function"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-7"
   },
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "def show_Q_function_progress(env_desc, V_all, pi_all):\n",
    "    # This function visualizes how value and policy changes over time.\n",
    "    # V: [num_iters, num_states]\n",
    "    # pi: [num_iters, num_states]\n",
    "\n",
    "    # We want to only shows few values\n",
    "    num_iters_all = V_all.shape[0]\n",
    "    num_iters = num_iters_all // 10\n",
    "\n",
    "    vis_indx = np.arange(0, num_iters_all, num_iters).tolist()\n",
    "    vis_indx.append(num_iters_all - 1)\n",
    "    V = np.zeros((len(vis_indx), V_all.shape[1]))\n",
    "    pi = np.zeros((len(vis_indx), V_all.shape[1]))\n",
    "\n",
    "    for c, i in enumerate(vis_indx):\n",
    "        V[c]  = V_all[i]\n",
    "        pi[c] = pi_all[i]\n",
    "\n",
    "    num_iters = V.shape[0]\n",
    "    fig, ax = plt.subplots(figsize=(15, 15))\n",
    "\n",
    "    for k in range(V.shape[0]):\n",
    "        plt.subplot(4, 4, k + 1)\n",
    "        plt.imshow(V[k].reshape(4,4), cmap=\"bone\")\n",
    "        ax = plt.gca()\n",
    "        ax.set_xticks(np.arange(0, 5)-.5, minor=True)\n",
    "        ax.set_yticks(np.arange(0, 5)-.5, minor=True)\n",
    "        ax.grid(which=\"minor\", color=\"w\", linestyle='-', linewidth=3)\n",
    "        ax.tick_params(which=\"minor\", bottom=False, left=False)\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])\n",
    "\n",
    "        # LEFT action: 0, DOWN action: 1\n",
    "        # RIGHT action: 2, UP action: 3\n",
    "        action2dxdy = {0:(-.25, 0),1:(0, .25),\n",
    "                       2:(0.25, 0),3:(-.25, 0)}\n",
    "\n",
    "        for y in range(4):\n",
    "            for x in range(4):\n",
    "                action = pi[k].reshape(4,4)[y, x]\n",
    "                dx, dy = action2dxdy[action]\n",
    "\n",
    "                if env_desc[y,x].decode() == 'H':\n",
    "                    ax.text(x, y, str(env_desc[y,x].decode()),\n",
    "                       ha=\"center\", va=\"center\", color=\"y\",\n",
    "                         size=20, fontweight='bold')\n",
    "\n",
    "                elif env_desc[y,x].decode() == 'G':\n",
    "                    ax.text(x, y, str(env_desc[y,x].decode()),\n",
    "                       ha=\"center\", va=\"center\", color=\"w\",\n",
    "                         size=20, fontweight='bold')\n",
    "\n",
    "                else:\n",
    "                    ax.text(x, y, str(env_desc[y,x].decode()),\n",
    "                       ha=\"center\", va=\"center\", color=\"g\",\n",
    "                         size=15, fontweight='bold')\n",
    "\n",
    "                # No arrow for cells with G and H labels\n",
    "                if env_desc[y,x].decode() != 'G' and env_desc[y,x].decode() != 'H':\n",
    "                    ax.arrow(x, y, dx, dy, color='r', head_width=0.2, head_length=0.15)\n",
    "\n",
    "        ax.set_title(\"Step = \"  + str(vis_indx[k] + 1), fontsize=20)\n",
    "\n",
    "    fig.tight_layout()\n",
    "    plt.show()"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-8 fw=all -->\n",
    "\n",
    "Trainer\n",
    "\n",
    "Legacy helper functions retained for backward compatibility:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-8"
   },
   "source": [
    "\n",
    "def load_array(data_arrays, batch_size, is_train=True):\n",
    "    \"\"\"Construct a PyTorch data iterator.\"\"\"\n",
    "    dataset = torch.utils.data.TensorDataset(*data_arrays)\n",
    "    return torch.utils.data.DataLoader(dataset, batch_size, shuffle=is_train)\n",
    "\n",
    "def synthetic_data(w, b, num_examples):\n",
    "    \"\"\"Generate y = Xw + b + noise.\"\"\"\n",
    "    X = d2l.normal(0, 1, (num_examples, len(w)))\n",
    "    y = d2l.matmul(X, w) + b\n",
    "    y += d2l.normal(0, 0.01, y.shape)\n",
    "    return X, d2l.reshape(y, (-1, 1))\n",
    "\n",
    "def sgd(params, lr, batch_size):\n",
    "    \"\"\"Minibatch stochastic gradient descent.\"\"\"\n",
    "    with torch.no_grad():\n",
    "        for param in params:\n",
    "            param -= lr * param.grad / batch_size\n",
    "            param.grad.zero_()\n",
    "\n",
    "def get_dataloader_workers():\n",
    "    \"\"\"Use 4 processes to read the data.\"\"\"\n",
    "    return 4\n",
    "\n",
    "def load_data_fashion_mnist(batch_size, resize=None):\n",
    "    \"\"\"Download the Fashion-MNIST dataset and then load it into memory.\"\"\"\n",
    "    trans = [transforms.ToTensor()]\n",
    "    if resize:\n",
    "        trans.insert(0, transforms.Resize(resize))\n",
    "    trans = transforms.Compose(trans)\n",
    "    mnist_train = torchvision.datasets.FashionMNIST(\n",
    "        root=\"../data\", train=True, transform=trans, download=True)\n",
    "    mnist_test = torchvision.datasets.FashionMNIST(\n",
    "        root=\"../data\", train=False, transform=trans, download=True)\n",
    "    return (torch.utils.data.DataLoader(mnist_train, batch_size, shuffle=True,\n",
    "                                        num_workers=get_dataloader_workers()),\n",
    "            torch.utils.data.DataLoader(mnist_test, batch_size, shuffle=False,\n",
    "                                        num_workers=get_dataloader_workers()))\n",
    "\n",
    "def evaluate_accuracy_gpu(net, data_iter, device=None):\n",
    "    \"\"\"Compute the accuracy for a model on a dataset using a GPU.\"\"\"\n",
    "    if isinstance(net, nn.Module):\n",
    "        net.eval()  # Set the model to evaluation mode\n",
    "        if not device:\n",
    "            device = next(iter(net.parameters())).device\n",
    "    # No. of correct predictions, no. of predictions\n",
    "    metric = d2l.Accumulator(2)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for X, y in data_iter:\n",
    "            if isinstance(X, list):\n",
    "                # Required for BERT Fine-tuning (to be covered later)\n",
    "                X = [x.to(device) for x in X]\n",
    "            else:\n",
    "                X = X.to(device)\n",
    "            y = y.to(device)\n",
    "            metric.add(d2l.accuracy(net(X), y), d2l.size(y))\n",
    "    return metric[0] / metric[1]\n",
    "\n",
    "\n",
    "def train_ch6(net, train_iter, test_iter, num_epochs, lr, device):\n",
    "    \"\"\"Train a model with a GPU (defined in Chapter 6).\"\"\"\n",
    "    def init_weights(m):\n",
    "        if type(m) == nn.Linear or type(m) == nn.Conv2d:\n",
    "            nn.init.xavier_uniform_(m.weight)\n",
    "    net.apply(init_weights)\n",
    "    print('training on', device)\n",
    "    net.to(device)\n",
    "    optimizer = torch.optim.SGD(net.parameters(), lr=lr)\n",
    "    loss = nn.CrossEntropyLoss()\n",
    "    animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs],\n",
    "                            legend=['train loss', 'train acc', 'test acc'])\n",
    "    timer, num_batches = d2l.Timer(), len(train_iter)\n",
    "    for epoch in range(num_epochs):\n",
    "        # Sum of training loss, sum of training accuracy, no. of examples\n",
    "        metric = d2l.Accumulator(3)\n",
    "        net.train()\n",
    "        for i, (X, y) in enumerate(train_iter):\n",
    "            timer.start()\n",
    "            optimizer.zero_grad()\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            y_hat = net(X)\n",
    "            l = loss(y_hat, y)\n",
    "            l.backward()\n",
    "            optimizer.step()\n",
    "            with torch.no_grad():\n",
    "                metric.add(l * X.shape[0], d2l.accuracy(y_hat, y), X.shape[0])\n",
    "            timer.stop()\n",
    "            train_l = metric[0] / metric[2]\n",
    "            train_acc = metric[1] / metric[2]\n",
    "            if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:\n",
    "                animator.add(epoch + (i + 1) / num_batches,\n",
    "                             (train_l, train_acc, None))\n",
    "        test_acc = evaluate_accuracy_gpu(net, test_iter)\n",
    "        animator.add(epoch + 1, (None, None, test_acc))\n",
    "    print(f'loss {train_l:.3f}, train acc {train_acc:.3f}, '\n",
    "          f'test acc {test_acc:.3f}')\n",
    "    print(f'{metric[2] * num_epochs / timer.sum():.1f} examples/sec '\n",
    "          f'on {str(device)}')"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-10"
   },
   "source": [
    "def show_images(imgs, num_rows, num_cols, titles=None, scale=1.5):\n",
    "    \"\"\"Plot a list of images.\"\"\"\n",
    "    figsize = (num_cols * scale, num_rows * scale)\n",
    "    _, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize)\n",
    "    axes = axes.flatten()\n",
    "    for i, (ax, img) in enumerate(zip(axes, imgs)):\n",
    "        try:\n",
    "            img = d2l.numpy(img)\n",
    "        except:\n",
    "            pass\n",
    "        ax.imshow(img)\n",
    "        ax.axes.get_xaxis().set_visible(False)\n",
    "        ax.axes.get_yaxis().set_visible(False)\n",
    "        if titles:\n",
    "            ax.set_title(titles[i])\n",
    "    return axes"
   ],
   "execution_count": 9,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-11"
   },
   "source": [
    "\n",
    "def linreg(X, w, b):\n",
    "    \"\"\"The linear regression model.\"\"\"\n",
    "    return d2l.matmul(X, w) + b\n",
    "\n",
    "def squared_loss(y_hat, y):\n",
    "    \"\"\"Squared loss.\"\"\"\n",
    "    return (y_hat - d2l.reshape(y, y_hat.shape)) ** 2 / 2\n",
    "\n",
    "def get_fashion_mnist_labels(labels):\n",
    "    \"\"\"Return text labels for the Fashion-MNIST dataset.\"\"\"\n",
    "    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat',\n",
    "                   'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']\n",
    "    return [text_labels[int(i)] for i in labels]\n",
    "\n",
    "#@tab pytorch, mxnet, tensorflow\n",
    "class Animator:\n",
    "    \"\"\"For plotting data in animation.\"\"\"\n",
    "    def __init__(self, xlabel=None, ylabel=None, legend=None, xlim=None,\n",
    "                 ylim=None, xscale='linear', yscale='linear',\n",
    "                 fmts=('-', 'm--', 'g-.', 'r:'), nrows=1, ncols=1,\n",
    "                 figsize=(3.5, 2.5)):\n",
    "        # Incrementally plot multiple lines\n",
    "        if legend is None:\n",
    "            legend = []\n",
    "        d2l.use_svg_display()\n",
    "        self.fig, self.axes = d2l.plt.subplots(nrows, ncols, figsize=figsize)\n",
    "        if nrows * ncols == 1:\n",
    "            self.axes = [self.axes, ]\n",
    "        # Use a lambda function to capture arguments\n",
    "        self.config_axes = lambda: d2l.set_axes(\n",
    "            self.axes[0], xlabel, ylabel, xlim, ylim, xscale, yscale, legend)\n",
    "        self.X, self.Y, self.fmts = None, None, fmts\n",
    "\n",
    "    def add(self, x, y):\n",
    "        # Add multiple data points into the figure\n",
    "        if not hasattr(y, \"__len__\"):\n",
    "            y = [y]\n",
    "        n = len(y)\n",
    "        if not hasattr(x, \"__len__\"):\n",
    "            x = [x] * n\n",
    "        if not self.X:\n",
    "            self.X = [[] for _ in range(n)]\n",
    "        if not self.Y:\n",
    "            self.Y = [[] for _ in range(n)]\n",
    "        for i, (a, b) in enumerate(zip(x, y)):\n",
    "            if a is not None and b is not None:\n",
    "                self.X[i].append(a)\n",
    "                self.Y[i].append(b)\n",
    "        self.axes[0].cla()\n",
    "        for x, y, fmt in zip(self.X, self.Y, self.fmts):\n",
    "            self.axes[0].plot(x, y, fmt)\n",
    "        self.config_axes()\n",
    "        display.display(self.fig)\n",
    "        display.clear_output(wait=True)\n",
    "\n",
    "#@tab pytorch, mxnet, tensorflow\n",
    "class Accumulator:\n",
    "    \"\"\"For accumulating sums over `n` variables.\"\"\"\n",
    "    def __init__(self, n):\n",
    "        self.data = [0.0] * n\n",
    "\n",
    "    def add(self, *args):\n",
    "        self.data = [a + float(b) for a, b in zip(self.data, args)]\n",
    "\n",
    "    def reset(self):\n",
    "        self.data = [0.0] * len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx]\n",
    "\n",
    "\n",
    "#@tab pytorch, mxnet, tensorflow\n",
    "def accuracy(y_hat, y):\n",
    "    \"\"\"Compute the number of correct predictions.\"\"\"\n",
    "    if len(y_hat.shape) > 1 and y_hat.shape[1] > 1:\n",
    "        y_hat = d2l.argmax(y_hat, axis=1)\n",
    "    elif (len(y_hat.shape) > 1 and y_hat.shape[-1] == 1) or y_hat.dtype != y.dtype:\n",
    "        # Binary classification with float scores (logits or probabilities):\n",
    "        # threshold at 0 (logits) to get class labels, then reshape to match y.\n",
    "        y_hat = d2l.astype(y_hat > 0, y.dtype).reshape(y.shape)\n",
    "    cmp = d2l.astype(y_hat, y.dtype) == y\n",
    "    return float(d2l.reduce_sum(d2l.astype(cmp, y.dtype)))"
   ],
   "execution_count": 10,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-12"
   },
   "source": [
    "\n",
    "import os\n",
    "import requests\n",
    "import zipfile\n",
    "import tarfile\n",
    "import hashlib\n",
    "\n",
    "def download(url, folder='../data', sha1_hash=None):\n",
    "    \"\"\"Download a file to folder and return the local filepath.\"\"\"\n",
    "    if not url.startswith('http'):\n",
    "        # For back compatibility\n",
    "        url, sha1_hash = DATA_HUB[url]\n",
    "    os.makedirs(folder, exist_ok=True)\n",
    "    fname = os.path.join(folder, url.split('/')[-1])\n",
    "    # Check if hit cache\n",
    "    if os.path.exists(fname) and sha1_hash:\n",
    "        sha1 = hashlib.sha1()\n",
    "        with open(fname, 'rb') as f:\n",
    "            while True:\n",
    "                data = f.read(1048576)\n",
    "                if not data:\n",
    "                    break\n",
    "                sha1.update(data)\n",
    "        if sha1.hexdigest() == sha1_hash:\n",
    "            return fname\n",
    "    # Download\n",
    "    import tempfile\n",
    "    print(f'Downloading {fname} from {url}...')\n",
    "    r = requests.get(url, stream=True, verify=True)\n",
    "    fd, tmp = tempfile.mkstemp(dir=folder)\n",
    "    try:\n",
    "        with os.fdopen(fd, 'wb') as f:\n",
    "            f.write(r.content)\n",
    "        os.replace(tmp, fname)\n",
    "    except BaseException:\n",
    "        os.unlink(tmp)\n",
    "        raise\n",
    "    return fname\n",
    "\n",
    "def extract(filename, folder=None):\n",
    "    \"\"\"Extract a zip/tar file into folder.\"\"\"\n",
    "    import tempfile, shutil\n",
    "    base_dir = os.path.dirname(filename)\n",
    "    _, ext = os.path.splitext(filename)\n",
    "    assert ext in ('.zip', '.tar', '.gz'), 'Only support zip/tar files.'\n",
    "    opener = zipfile.ZipFile if ext == '.zip' else tarfile.open\n",
    "    if folder is None:\n",
    "        folder = base_dir\n",
    "    tmp = tempfile.mkdtemp(dir=folder)\n",
    "    try:\n",
    "        with opener(filename, 'r') as fp:\n",
    "            fp.extractall(tmp)\n",
    "        for name in os.listdir(tmp):\n",
    "            src = os.path.join(tmp, name)\n",
    "            dst = os.path.join(folder, name)\n",
    "            if os.path.isdir(dst):\n",
    "                shutil.rmtree(dst)\n",
    "            os.rename(src, dst)\n",
    "    finally:\n",
    "        shutil.rmtree(tmp, ignore_errors=True)"
   ],
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-13"
   },
   "source": [
    "\n",
    "def download_extract(name, folder=None):\n",
    "    \"\"\"Download and extract a zip/tar file.\"\"\"\n",
    "    fname = download(name)\n",
    "    base_dir = os.path.dirname(fname)\n",
    "    data_dir, ext = os.path.splitext(fname)\n",
    "    target = os.path.join(base_dir, folder) if folder else data_dir\n",
    "    # Skip re-extraction if a completion marker exists (extracting many small\n",
    "    # files is slow and unnecessary when the archive is already unpacked).\n",
    "    marker = fname + '.extracted'\n",
    "    if os.path.exists(marker):\n",
    "        return target\n",
    "    if ext == '.zip':\n",
    "        fp = zipfile.ZipFile(fname, 'r')\n",
    "    elif ext in ('.tar', '.gz'):\n",
    "        fp = tarfile.open(fname, 'r')\n",
    "    else:\n",
    "        assert False, 'Only zip/tar files can be extracted.'\n",
    "    fp.extractall(base_dir)\n",
    "    open(marker, 'w').close()\n",
    "    return target\n",
    "\n",
    "\n",
    "def tokenize(lines, token='word'):\n",
    "    \"\"\"Split text lines into word or character tokens.\"\"\"\n",
    "    assert token in ('word', 'char'), 'Unknown token type: ' + token\n",
    "    return [line.split() if token == 'word' else list(line) for line in lines]"
   ],
   "execution_count": 12,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-14"
   },
   "source": [
    "\n",
    "def evaluate_loss(net, data_iter, loss):\n",
    "    \"\"\"Evaluate the loss of a model on the given dataset.\"\"\"\n",
    "    metric = d2l.Accumulator(2)  # Sum of losses, no. of examples\n",
    "    for X, y in data_iter:\n",
    "        out = net(X)\n",
    "        y = d2l.reshape(y, out.shape)\n",
    "        l = loss(out, y)\n",
    "        metric.add(d2l.reduce_sum(l), d2l.size(l))\n",
    "    return metric[0] / metric[1]"
   ],
   "execution_count": 13,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-15"
   },
   "source": [
    "def grad_clipping(net, theta):\n",
    "    \"\"\"Clip the gradient.\"\"\"\n",
    "    if isinstance(net, nn.Module):\n",
    "        params = [p for p in net.parameters() if p.requires_grad]\n",
    "    else:\n",
    "        params = net.params\n",
    "    norm = torch.sqrt(sum(torch.sum((p.grad ** 2)) for p in params))\n",
    "    if norm > theta:\n",
    "        for param in params:\n",
    "            param.grad[:] *= theta / norm"
   ],
   "execution_count": 14,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-9 fw=all -->\n",
    "\n",
    "More for the attention chapter."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-16"
   },
   "source": [
    "\n",
    "d2l.DATA_HUB['fra-eng'] = (d2l.DATA_URL + 'fra-eng.zip',\n",
    "                           '94646ad1522d915e7b0f9296181140edcf86a4f5')\n",
    "\n",
    "\n",
    "def read_data_nmt():\n",
    "    \"\"\"Load the English-French dataset.\"\"\"\n",
    "    data_dir = d2l.download_extract('fra-eng')\n",
    "    with open(os.path.join(data_dir, 'fra.txt'), 'r', encoding='utf-8') as f:\n",
    "        return f.read()\n",
    "\n",
    "\n",
    "def preprocess_nmt(text):\n",
    "    \"\"\"Preprocess the English-French dataset.\"\"\"\n",
    "    def no_space(char, prev_char):\n",
    "        return char in set(',.!?') and prev_char != ' '\n",
    "\n",
    "    # Replace non-breaking space with space, and convert uppercase letters to\n",
    "    # lowercase ones\n",
    "    text = text.replace('\\u202f', ' ').replace('\\xa0', ' ').lower()\n",
    "    # Insert space between words and punctuation marks\n",
    "    out = [' ' + char if i > 0 and no_space(char, text[i - 1]) else char\n",
    "           for i, char in enumerate(text)]\n",
    "    return ''.join(out)\n",
    "\n",
    "\n",
    "def tokenize_nmt(text, num_examples=None):\n",
    "    \"\"\"Tokenize the English-French dataset.\"\"\"\n",
    "    source, target = [], []\n",
    "    for i, line in enumerate(text.split('\\n')):\n",
    "        if num_examples and i > num_examples:\n",
    "            break\n",
    "        parts = line.split('\\t')\n",
    "        if len(parts) == 2:\n",
    "            source.append(parts[0].split(' '))\n",
    "            target.append(parts[1].split(' '))\n",
    "    return source, target\n",
    "\n",
    "    \n",
    "\n",
    "def truncate_pad(line, num_steps, padding_token):\n",
    "    \"\"\"Truncate or pad sequences.\"\"\"\n",
    "    if len(line) > num_steps:\n",
    "        return line[:num_steps]  # Truncate\n",
    "    return line + [padding_token] * (num_steps - len(line))  # Pad\n",
    "\n",
    "\n",
    "def build_array_nmt(lines, vocab, num_steps):\n",
    "    \"\"\"Transform text sequences of machine translation into minibatches.\"\"\"\n",
    "    lines = [vocab[l] for l in lines]\n",
    "    lines = [l + [vocab['<eos>']] for l in lines]\n",
    "    array = d2l.tensor([truncate_pad(\n",
    "        l, num_steps, vocab['<pad>']) for l in lines])\n",
    "    valid_len = d2l.reduce_sum(\n",
    "        d2l.astype(array != vocab['<pad>'], d2l.int32), 1)\n",
    "    return array, valid_len\n",
    "\n",
    "\n",
    "def load_data_nmt(batch_size, num_steps, num_examples=600):\n",
    "    \"\"\"Return the iterator and the vocabularies of the translation dataset.\"\"\"\n",
    "    text = preprocess_nmt(read_data_nmt())\n",
    "    source, target = tokenize_nmt(text, num_examples)\n",
    "    src_vocab = d2l.Vocab(source, min_freq=2,\n",
    "                          reserved_tokens=['<pad>', '<bos>', '<eos>'])\n",
    "    tgt_vocab = d2l.Vocab(target, min_freq=2,\n",
    "                          reserved_tokens=['<pad>', '<bos>', '<eos>'])\n",
    "    src_array, src_valid_len = build_array_nmt(source, src_vocab, num_steps)\n",
    "    tgt_array, tgt_valid_len = build_array_nmt(target, tgt_vocab, num_steps)\n",
    "    data_arrays = (src_array, src_valid_len, tgt_array, tgt_valid_len)\n",
    "    data_iter = d2l.load_array(data_arrays, batch_size)\n",
    "    return data_iter, src_vocab, tgt_vocab"
   ],
   "execution_count": 15,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-17"
   },
   "source": [
    "\n",
    "def sequence_mask(X, valid_len, value=0):\n",
    "    \"\"\"Mask irrelevant entries in sequences.\"\"\"\n",
    "    maxlen = X.size(1)\n",
    "    mask = torch.arange((maxlen), dtype=torch.float32,\n",
    "                        device=X.device)[None, :] < valid_len[:, None]\n",
    "    X[~mask] = value\n",
    "    return X\n",
    "\n",
    "    \n",
    "\n",
    "class MaskedSoftmaxCELoss(nn.CrossEntropyLoss):\n",
    "    \"\"\"The softmax cross-entropy loss with masks.\"\"\"\n",
    "    # `pred` shape: (`batch_size`, `num_steps`, `vocab_size`)\n",
    "    # `label` shape: (`batch_size`, `num_steps`)\n",
    "    # `valid_len` shape: (`batch_size`,)\n",
    "    def __init__(self):\n",
    "        super().__init__(reduction='none')\n",
    "\n",
    "    def forward(self, pred, label, valid_len):\n",
    "        weights = torch.ones_like(label)\n",
    "        weights = sequence_mask(weights, valid_len)\n",
    "        unweighted_loss = super().forward(pred.permute(0, 2, 1), label)\n",
    "        weighted_loss = (unweighted_loss * weights).mean(dim=1)\n",
    "        return weighted_loss\n",
    "    \n",
    "\n",
    "def train_seq2seq(net, data_iter, lr, num_epochs, tgt_vocab, device):\n",
    "    \"\"\"Train a model for sequence to sequence.\"\"\"\n",
    "    def xavier_init_weights(m):\n",
    "        if type(m) == nn.Linear:\n",
    "            nn.init.xavier_uniform_(m.weight)\n",
    "        if type(m) == nn.GRU:\n",
    "            for param in m._flat_weights_names:\n",
    "                if \"weight\" in param:\n",
    "                    nn.init.xavier_uniform_(m._parameters[param])\n",
    "    net.apply(xavier_init_weights)\n",
    "    net.to(device)\n",
    "    optimizer = torch.optim.Adam(net.parameters(), lr=lr)\n",
    "    loss = MaskedSoftmaxCELoss()\n",
    "    net.train()\n",
    "    animator = d2l.Animator(xlabel='epoch', ylabel='loss',\n",
    "                            xlim=[10, num_epochs])\n",
    "    for epoch in range(num_epochs):\n",
    "        timer = d2l.Timer()\n",
    "        metric = d2l.Accumulator(2)  # Sum of training loss, no. of tokens\n",
    "        for batch in data_iter:\n",
    "            optimizer.zero_grad()\n",
    "            X, X_valid_len, Y, Y_valid_len = [x.to(device) for x in batch]\n",
    "            bos = torch.tensor([tgt_vocab['<bos>']] * Y.shape[0],\n",
    "                               device=device).reshape(-1, 1)\n",
    "            dec_input = d2l.concat([bos, Y[:, :-1]], 1)  # Teacher forcing\n",
    "            Y_hat, _ = net(X, dec_input, X_valid_len)\n",
    "            l = loss(Y_hat, Y, Y_valid_len)\n",
    "            l.sum().backward()  # Make the loss scalar for `backward`\n",
    "            d2l.grad_clipping(net, 1)\n",
    "            num_tokens = Y_valid_len.sum()\n",
    "            optimizer.step()\n",
    "            with torch.no_grad():\n",
    "                metric.add(l.sum(), num_tokens)\n",
    "        if (epoch + 1) % 10 == 0:\n",
    "            animator.add(epoch + 1, (metric[0] / metric[1],))\n",
    "    print(f'loss {metric[0] / metric[1]:.3f}, {metric[1] / timer.stop():.1f} '\n",
    "          f'tokens/sec on {str(device)}')\n",
    "    \n",
    "\n",
    "\n",
    "def predict_seq2seq(net, src_sentence, src_vocab, tgt_vocab, num_steps,\n",
    "                    device, save_attention_weights=False):\n",
    "    \"\"\"Predict for sequence to sequence.\"\"\"\n",
    "    # Set `net` to eval mode for inference\n",
    "    net.eval()\n",
    "    src_tokens = src_vocab[src_sentence.lower().split(' ')] + [\n",
    "        src_vocab['<eos>']]\n",
    "    enc_valid_len = torch.tensor([len(src_tokens)], device=device)\n",
    "    src_tokens = d2l.truncate_pad(src_tokens, num_steps, src_vocab['<pad>'])\n",
    "    # Add the batch axis\n",
    "    enc_X = torch.unsqueeze(\n",
    "        torch.tensor(src_tokens, dtype=torch.long, device=device), dim=0)\n",
    "    enc_outputs = net.encoder(enc_X, enc_valid_len)\n",
    "    dec_state = net.decoder.init_state(enc_outputs, enc_valid_len)\n",
    "    # Add the batch axis\n",
    "    dec_X = torch.unsqueeze(torch.tensor(\n",
    "        [tgt_vocab['<bos>']], dtype=torch.long, device=device), dim=0)\n",
    "    output_seq, attention_weight_seq = [], []\n",
    "    for _ in range(num_steps):\n",
    "        Y, dec_state = net.decoder(dec_X, dec_state)\n",
    "        # We use the token with the highest prediction likelihood as input\n",
    "        # of the decoder at the next time step\n",
    "        dec_X = Y.argmax(dim=2)\n",
    "        pred = dec_X.squeeze(dim=0).type(torch.int32).item()\n",
    "        # Save attention weights (to be covered later)\n",
    "        if save_attention_weights:\n",
    "            attention_weight_seq.append(net.decoder.attention_weights)\n",
    "        # Once the end-of-sequence token is predicted, the generation of the\n",
    "        # output sequence is complete\n",
    "        if pred == tgt_vocab['<eos>']:\n",
    "            break\n",
    "        output_seq.append(pred)\n",
    "    return ' '.join(tgt_vocab.to_tokens(output_seq)), attention_weight_seq"
   ],
   "execution_count": 16,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-10 fw=all -->\n",
    "\n",
    "## A Note on Framework Coverage\n",
    "\n",
    "The legacy helpers in this section (`evaluate_accuracy`, `train_ch6`,\n",
    "`train_seq2seq`, `predict_seq2seq`, `MaskedSoftmaxCELoss`) are kept for\n",
    "parity with the original D2L implementation, which predates the unified\n",
    "`d2l.Trainer` class introduced in this edition. They are deliberately\n",
    "not provided for JAX, and PyTorch only ships the subset that is genuinely\n",
    "useful outside of the `Trainer` flow. If you are reading the JAX tab,\n",
    "the corresponding chapters use `d2l.Trainer.fit(model, data)` end-to-end;\n",
    "the per-batch logic that these helpers spell out lives inside\n",
    "`Trainer.fit_epoch` and the `@d2l.add_to_class` extensions defined\n",
    "alongside it. The MXNet and TensorFlow tabs additionally retain\n",
    "`evaluate_accuracy` because some of their earlier-chapter snippets call\n",
    "it directly; for PyTorch and JAX,\n",
    "use `Trainer.validation_step` (or its return value) instead."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
