{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-1 fw=all -->\n",
    "\n",
    "# Residual Networks: ResNet, ResNeXt, and DenseNet\n",
    "\n",
    "As we design ever deeper networks it becomes imperative to understand how adding layers can increase the complexity and expressiveness of the network.\n",
    "Even more important is the ability to design networks where adding layers\n",
    "preserves every function the shallower network could represent, while making\n",
    "additional functions available.\n",
    "To make some progress we need a bit of mathematics."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-residual-networks-resnet-and-resnext"
   },
   "source": [
    "from d2l import jax as d2l\n",
    "from flax import nnx\n",
    "from jax import numpy as jnp\n",
    "import jax"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-2 fw=all -->\n",
    "\n",
    "## Function Classes\n",
    "\n",
    "Consider $\\mathcal{F}$, the class of functions represented by a network\n",
    "architecture as its weights and biases vary. This definition concerns the\n",
    "parameterization; which members an optimizer can reach is a separate question.\n",
    "Let's assume that $f^*$ is the \"truth\" function that we really would like to find.\n",
    "If it is in $\\mathcal{F}$, we are in good shape but typically we will not be quite so lucky.\n",
    "Instead, we will try to find some $f^*_{\\mathcal{F}}$ which is our best bet within $\\mathcal{F}$.\n",
    "For instance,\n",
    "given a dataset with features $\\mathbf{X}$\n",
    "and labels $\\mathbf{y}$,\n",
    "we might try finding it by solving the following optimization problem:\n",
    "\n",
    "$$f^*_{\\mathcal{F}} \\stackrel{\\textrm{def}}{=} \\mathop{\\mathrm{argmin}}_f L(\\mathbf{X}, \\mathbf{y}, f) \\textrm{ subject to } f \\in \\mathcal{F}.$$\n",
    "\n",
    "We know that regularization [@tikhonov1977solutions; @morozov2012methods] may control complexity of $\\mathcal{F}$\n",
    "and achieve consistency, so a larger size of training data\n",
    "generally leads to better $f^*_{\\mathcal{F}}$.\n",
    "It is only reasonable to assume that if we design a different and more powerful architecture $\\mathcal{F}'$ we should arrive at a better outcome. In other words, we would expect that $f^*_{\\mathcal{F}'}$ is \"better\" than $f^*_{\\mathcal{F}}$. However, if $\\mathcal{F} \\not\\subseteq \\mathcal{F}'$ there is no guarantee that this should even happen. In fact, $f^*_{\\mathcal{F}'}$ might well be worse.\n",
    "As illustrated by the figure,\n",
    "for non-nested function classes, a larger function class does not always move closer to the \"truth\" function $f^*$. For instance,\n",
    "on the left of the figure,\n",
    "though $\\mathcal{F}_3$ is closer to $f^*$ than $\\mathcal{F}_1$, $\\mathcal{F}_6$ moves away and there is no guarantee that further increasing the complexity can reduce the distance from $f^*$.\n",
    "With nested function classes\n",
    "where $\\mathcal{F}_1 \\subseteq \\cdots \\subseteq \\mathcal{F}_6$\n",
    "on the right of the figure,\n",
    "we can avoid the aforementioned issue from the non-nested function classes.\n",
    "\n",
    "\n",
    "![For non-nested function classes, a larger (indicated by area) function class does not guarantee we will get closer to the \"truth\" function ($\\mathit{f}^*$). This does not happen in nested function classes.](../img/functionclasses.svg)\n",
    "\n",
    "Thus, if larger function classes contain the smaller ones, increasing capacity\n",
    "cannot reduce expressive power. The containment may be equality, so the\n",
    "increase need not be strict.\n",
    "For deep neural networks,\n",
    "if a newly added block can represent the identity\n",
    "$f(\\mathbf{x}) = \\mathbf{x}$, the deeper model contains the shallower one as a\n",
    "special case. This protects representation capacity; it does not guarantee\n",
    "that optimization will find the best member of the larger class.\n",
    "\n",
    "This is the question that @He.Zhang.Ren.ea.2016 considered when working on very deep computer vision models.\n",
    "At the heart of their proposed *residual network* (*ResNet*) is the idea that every additional layer should\n",
    "more easily\n",
    "contain the identity function as one of its elements.\n",
    "These considerations led to a simple solution, a *residual block*.\n",
    "With it, ResNet won the ImageNet Large Scale Visual Recognition Challenge in 2015. The design influenced how to\n",
    "build deep neural networks. For instance, residual blocks have been added to recurrent networks [@prakash2016neural; @kim2017residual]. Likewise, Transformers [@Vaswani.Shazeer.Parmar.ea.2017] use them to stack many layers of networks efficiently. It is also used in graph neural networks [@Kipf.Welling.2016] and, as a basic concept, it has been used extensively in computer vision [@Redmon.Farhadi.2018; @Ren.He.Girshick.ea.2015]. \n",
    "Highway networks [@srivastava2015highway] predate ResNet and share some\n",
    "of its motivation, using learned gates rather than a direct identity shortcut.\n",
    "\n",
    "\n",
    "## Residual Blocks\n",
    "\n",
    "Let's focus on a local part of a neural network, as depicted in the figure. Denote the input by $\\mathbf{x}$.\n",
    "We assume that $f(\\mathbf{x})$, the desired underlying mapping we want to obtain by learning, is to be used as input to the activation function on the top.\n",
    "On the left,\n",
    "the portion within the dotted-line box\n",
    "must directly learn $f(\\mathbf{x})$.\n",
    "On the right,\n",
    "the portion within the dotted-line box\n",
    "needs to\n",
    "learn the *residual mapping* $g(\\mathbf{x}) = f(\\mathbf{x}) - \\mathbf{x}$,\n",
    "which is how the residual block derives its name.\n",
    "If the identity mapping $f(\\mathbf{x}) = \\mathbf{x}$ is the desired underlying mapping,\n",
    "the residual mapping amounts to $g(\\mathbf{x}) = 0$ and it is thus easier to learn:\n",
    "we only need to push the weights and biases\n",
    "of the\n",
    "upper weight layer (e.g., fully connected layer and convolutional layer)\n",
    "within the dotted-line box\n",
    "to zero.\n",
    "The right figure illustrates the *residual block* of ResNet,\n",
    "where the solid line carrying the layer input\n",
    "$\\mathbf{x}$ to the addition operator\n",
    "is called a *residual connection* (or *shortcut connection*).\n",
    "Residual connections provide a direct path for activations and gradients\n",
    "through a stack of blocks.\n",
    "In fact,\n",
    "the residual block\n",
    "can be thought of as\n",
    "a special case of the multi-branch Inception block:\n",
    "it has two branches\n",
    "one of which is the identity mapping.\n",
    "\n",
    "![In a regular block (left), the portion within the dotted-line box must directly learn the mapping $\\mathit{f}(\\mathbf{x})$. In a residual block (right), the portion within the dotted-line box needs to learn the residual mapping $\\mathit{g}(\\mathbf{x}) = \\mathit{f}(\\mathbf{x}) - \\mathbf{x}$, making the identity mapping $\\mathit{f}(\\mathbf{x}) = \\mathbf{x}$ easier to learn.](../img/residual-block.svg)\n",
    "\n",
    "\n",
    "ResNet has VGG's full $3\\times 3$ convolutional layer design. The residual\n",
    "block has two $3\\times 3$ convolutional layers with the same number of output\n",
    "channels. Each convolutional layer is followed by batch normalization, and the\n",
    "first is followed by a ReLU. We add the input to the second normalized output\n",
    "and then apply the final ReLU. Because that ReLU clips negative values, this\n",
    "post-activation block is exactly the identity only on inputs that are already\n",
    "nonnegative. The pre-activation variant discussed in the exercises places the\n",
    "activation inside the residual branch and permits an exact identity on\n",
    "arbitrary inputs.\n",
    "This kind of design requires that the output of the two convolutional layers has to be of the same shape as the input, so that they can be added together. If we want to change the number of channels, we need to introduce an additional $1\\times 1$ convolutional layer to transform the input into the desired shape for the addition operation. Let's have a look at the code below."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-residual-blocks-1"
   },
   "source": [
    "class Residual(nnx.Module):\n",
    "    \"\"\"The Residual block of ResNet models.\"\"\"\n",
    "    def __init__(self, num_channels, use_1x1conv=False, strides=(1, 1),\n",
    "                 in_channels=None, rngs=None):\n",
    "        in_channels = num_channels if in_channels is None else in_channels\n",
    "        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs\n",
    "        self.conv1 = nnx.Conv(in_channels, num_channels, kernel_size=(3, 3),\n",
    "                              padding='same', strides=strides, rngs=rngs)\n",
    "        self.conv2 = nnx.Conv(num_channels, num_channels, kernel_size=(3, 3),\n",
    "                              padding='same', rngs=rngs)\n",
    "        # Auto-enable 1x1 conv when downsampling so the residual shape matches.\n",
    "        if use_1x1conv or any(s != 1 for s in strides):\n",
    "            self.conv3 = nnx.Conv(in_channels, num_channels,\n",
    "                                  kernel_size=(1, 1), strides=strides,\n",
    "                                  rngs=rngs)\n",
    "        else:\n",
    "            self.conv3 = None\n",
    "        self.bn1 = nnx.BatchNorm(num_channels, rngs=rngs)\n",
    "        self.bn2 = nnx.BatchNorm(num_channels, rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        Y = nnx.relu(self.bn1(self.conv1(X)))\n",
    "        Y = self.bn2(self.conv2(Y))\n",
    "        if self.conv3:\n",
    "            X = self.conv3(X)\n",
    "        Y += X\n",
    "        return nnx.relu(Y)"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-3 fw=all -->\n",
    "\n",
    "This code generates two types of networks: one where we add the input directly to the output before applying the ReLU nonlinearity whenever `use_1x1conv=False` and the stride is `1`; and one where we adjust channels and resolution by means of a $1 \\times 1$ convolution before adding, which the block enables automatically whenever `use_1x1conv=True` or the stride is not `1`. the figure illustrates this.\n",
    "\n",
    "![The two ResNet block variants side by side: with an identity skip when input and output shapes match (left), and with a $1 \\times 1$ convolution on the skip path that adjusts channels and resolution so the addition is well defined (right).](../img/arch-resnet-block.svg)\n",
    "\n",
    "Now let's look at a situation where the input and output are of the same shape, where $1 \\times 1$ convolution is not needed."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-residual-blocks-2"
   },
   "source": [
    "blk = Residual(3)\n",
    "X = jax.random.normal(d2l.get_key(), (4, 6, 6, 3))\n",
    "blk(X).shape"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(4, 6, 6, 3)"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-4 fw=all -->\n",
    "\n",
    "We also have the option to halve the output height and width while increasing the number of output channels.\n",
    "In this case we use $1 \\times 1$ convolutions via `use_1x1conv=True`. This comes in handy at the beginning of each ResNet block to reduce the spatial dimensionality via `strides=2`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-residual-blocks-3"
   },
   "source": [
    "blk = Residual(6, use_1x1conv=True, strides=(2, 2), in_channels=3)\n",
    "blk(X).shape"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(4, 3, 3, 6)"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-5 fw=all -->\n",
    "\n",
    "## ResNet Model\n",
    "\n",
    "The first two layers of ResNet are the same as those of the GoogLeNet we described before: the $7\\times 7$ convolutional layer with 64 output channels and a stride of 2 is followed by the $3\\times 3$ max-pooling layer with a stride of 2. The difference is the batch normalization layer added after each convolutional layer in ResNet."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-resnet-model-1"
   },
   "source": [
    "class ResNet(d2l.Classifier):\n",
    "    def __init__(self, arch, lr=0.1, num_classes=10, in_channels=1, rngs=None):\n",
    "        super().__init__()\n",
    "        self.save_hyperparameters(ignore=['rngs'])\n",
    "        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs\n",
    "        self.net = self.create_net(in_channels, rngs)\n",
    "\n",
    "    def b1(self, in_channels, rngs):\n",
    "        return nnx.Sequential(\n",
    "            nnx.Conv(in_channels, 64, kernel_size=(7, 7), strides=(2, 2),\n",
    "                     padding='same', rngs=rngs),\n",
    "            nnx.BatchNorm(64, rngs=rngs), nnx.relu,\n",
    "            lambda x: nnx.max_pool(x, window_shape=(3, 3), strides=(2, 2),\n",
    "                                   padding='same'))"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-6 fw=all -->\n",
    "\n",
    "GoogLeNet uses four modules made up of Inception blocks.\n",
    "However, ResNet uses four modules made up of residual blocks, each of which uses several residual blocks with the same number of output channels.\n",
    "The number of channels in the first module is the same as the number of input channels. Since a max-pooling layer with a stride of 2 has already been used, it is not necessary to reduce the height and width. In the first residual block for each of the subsequent modules, the number of channels is doubled compared with that of the previous module, and the height and width are halved."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-resnet-model-2"
   },
   "source": [
    "@d2l.add_to_class(ResNet)\n",
    "def block(self, num_residuals, num_channels, in_channels,\n",
    "          first_block=False, rngs=None):\n",
    "    blk = []\n",
    "    for i in range(num_residuals):\n",
    "        if i == 0 and not first_block:\n",
    "            blk.append(Residual(num_channels, use_1x1conv=True,\n",
    "                                strides=(2, 2), in_channels=in_channels,\n",
    "                                rngs=rngs))\n",
    "        else:\n",
    "            blk.append(Residual(num_channels, in_channels=in_channels,\n",
    "                                rngs=rngs))\n",
    "        in_channels = num_channels\n",
    "    return nnx.Sequential(*blk)"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-7 fw=all -->\n",
    "\n",
    "Then, we add all the modules to ResNet. Here, two residual blocks are used for each module. Lastly, just like GoogLeNet, we add a global average pooling layer, followed by the fully connected layer output."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-resnet-model-3"
   },
   "source": [
    "@d2l.add_to_class(ResNet)\n",
    "def create_net(self, in_channels, rngs):\n",
    "    layers = [self.b1(in_channels, rngs)]\n",
    "    stage_channels = 64\n",
    "    for i, (num_residuals, num_channels) in enumerate(self.arch):\n",
    "        layers.append(self.block(num_residuals, num_channels, stage_channels,\n",
    "                                 first_block=(i == 0), rngs=rngs))\n",
    "        stage_channels = num_channels\n",
    "    layers.append(nnx.Sequential(\n",
    "        lambda x: x.mean(axis=(1, 2)),  # global avg pooling over H, W (NHWC)\n",
    "        nnx.Linear(stage_channels, self.num_classes, rngs=rngs)))\n",
    "    return nnx.Sequential(*layers)"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-8 fw=all -->\n",
    "\n",
    "There are four convolutional layers in each module (excluding the $1\\times 1$ convolutional layer). Together with the first $7\\times 7$ convolutional layer and the final fully connected layer, there are 18 layers in total. Therefore, this model is commonly known as ResNet-18.\n",
    "By configuring different numbers of channels and residual blocks in the module, we can create different ResNet models, such as the deeper 152-layer ResNet-152. Although the main architecture of ResNet is similar to that of GoogLeNet, ResNet's structure is simpler and easier to modify. All these factors have resulted in the rapid and widespread use of ResNet. the figure depicts the full ResNet-18.\n",
    "\n",
    "![The full ResNet-18: a $7\\times 7$ convolutional stem, four stages of two residual blocks each, where the first block of stages 2--4 halves the resolution and doubles the channels, and a global average pooling head.](../img/arch-resnet18.svg)\n",
    "\n",
    "Before training ResNet, let's observe how the input shape changes across different modules in ResNet. As in all the previous architectures, the resolution decreases while the number of channels increases up until the point where a global average pooling layer aggregates all features."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-resnet-model-4"
   },
   "source": [
    "class ResNet18(ResNet):\n",
    "    def __init__(self, lr=0.1, num_classes=10, in_channels=1, rngs=None):\n",
    "        super().__init__(((2, 64), (2, 128), (2, 256), (2, 512)),\n",
    "                         lr, num_classes, in_channels, rngs)"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-resnet-model-5"
   },
   "source": [
    "ResNet18().layer_summary((1, 96, 96, 1))"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Sequential output shape:\t (1, 24, 24, 64)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Sequential output shape:\t (1, 24, 24, 64)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Sequential output shape:\t (1, 12, 12, 128)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Sequential output shape:\t (1, 6, 6, 256)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Sequential output shape:\t (1, 3, 3, 512)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Sequential output shape:\t (1, 10)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-9 fw=all -->\n",
    "\n",
    "## Training\n",
    "\n",
    "We train ResNet on the Fashion-MNIST dataset, just like before. ResNet is quite a powerful and flexible architecture. The plot capturing training and validation loss illustrates a significant gap between both graphs, with the training loss being considerably lower. For a network of this flexibility, more training data would offer distinct benefit in closing the gap and improving accuracy."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-training"
   },
   "source": [
    "model = ResNet18(lr=0.01)\n",
    "trainer = d2l.Trainer(max_epochs=10, num_gpus=1)\n",
    "data = d2l.FashionMNIST(batch_size=128, resize=(96, 96))\n",
    "trainer.fit(model, data)"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"238.965625pt\" height=\"183.35625pt\" viewBox=\"0 0 238.965625 183.35625\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-19T14:30:32.668097</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 183.35625 \nL 238.965625 183.35625 \nL 238.965625 0 \nL 0 0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 30.103125 145.8 \nL 225.403125 145.8 \nL 225.403125 7.2 \nL 30.103125 7.2 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <defs>\n       <path id=\"m3d5441be01\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m3d5441be01\" x=\"30.103125\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 0 -->\n      <g transform=\"translate(26.921875 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_2\">\n      <g>\n       <use xlink:href=\"#m3d5441be01\" x=\"69.163125\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 2 -->\n      <g transform=\"translate(65.981875 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-32\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_3\">\n      <g>\n       <use xlink:href=\"#m3d5441be01\" x=\"108.223125\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 4 -->\n      <g transform=\"translate(105.041875 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-34\" d=\"M 2419 4116 \nL 825 1625 \nL 2419 1625 \nL 2419 4116 \nz\nM 2253 4666 \nL 3047 4666 \nL 3047 1625 \nL 3713 1625 \nL 3713 1100 \nL 3047 1100 \nL 3047 0 \nL 2419 0 \nL 2419 1100 \nL 313 1100 \nL 313 1709 \nL 2253 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-34\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_4\">\n      <g>\n       <use xlink:href=\"#m3d5441be01\" x=\"147.283125\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 6 -->\n      <g transform=\"translate(144.101875 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-36\" d=\"M 2113 2584 \nQ 1688 2584 1439 2293 \nQ 1191 2003 1191 1497 \nQ 1191 994 1439 701 \nQ 1688 409 2113 409 \nQ 2538 409 2786 701 \nQ 3034 994 3034 1497 \nQ 3034 2003 2786 2293 \nQ 2538 2584 2113 2584 \nz\nM 3366 4563 \nL 3366 3988 \nQ 3128 4100 2886 4159 \nQ 2644 4219 2406 4219 \nQ 1781 4219 1451 3797 \nQ 1122 3375 1075 2522 \nQ 1259 2794 1537 2939 \nQ 1816 3084 2150 3084 \nQ 2853 3084 3261 2657 \nQ 3669 2231 3669 1497 \nQ 3669 778 3244 343 \nQ 2819 -91 2113 -91 \nQ 1303 -91 875 529 \nQ 447 1150 447 2328 \nQ 447 3434 972 4092 \nQ 1497 4750 2381 4750 \nQ 2619 4750 2861 4703 \nQ 3103 4656 3366 4563 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-36\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_5\">\n     <g id=\"line2d_5\">\n      <g>\n       <use xlink:href=\"#m3d5441be01\" x=\"186.343125\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_5\">\n      <!-- 8 -->\n      <g transform=\"translate(183.161875 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-38\" d=\"M 2034 2216 \nQ 1584 2216 1326 1975 \nQ 1069 1734 1069 1313 \nQ 1069 891 1326 650 \nQ 1584 409 2034 409 \nQ 2484 409 2743 651 \nQ 3003 894 3003 1313 \nQ 3003 1734 2745 1975 \nQ 2488 2216 2034 2216 \nz\nM 1403 2484 \nQ 997 2584 770 2862 \nQ 544 3141 544 3541 \nQ 544 4100 942 4425 \nQ 1341 4750 2034 4750 \nQ 2731 4750 3128 4425 \nQ 3525 4100 3525 3541 \nQ 3525 3141 3298 2862 \nQ 3072 2584 2669 2484 \nQ 3125 2378 3379 2068 \nQ 3634 1759 3634 1313 \nQ 3634 634 3220 271 \nQ 2806 -91 2034 -91 \nQ 1263 -91 848 271 \nQ 434 634 434 1313 \nQ 434 1759 690 2068 \nQ 947 2378 1403 2484 \nz\nM 1172 3481 \nQ 1172 3119 1398 2916 \nQ 1625 2713 2034 2713 \nQ 2441 2713 2670 2916 \nQ 2900 3119 2900 3481 \nQ 2900 3844 2670 4047 \nQ 2441 4250 2034 4250 \nQ 1625 4250 1398 4047 \nQ 1172 3844 1172 3481 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-38\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_6\">\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#m3d5441be01\" x=\"225.403125\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- 10 -->\n      <g transform=\"translate(219.040625 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_7\">\n     <!-- epoch -->\n     <g transform=\"translate(112.525 174.076563) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-70\" d=\"M 1159 525 \nL 1159 -1331 \nL 581 -1331 \nL 581 3500 \nL 1159 3500 \nL 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nz\nM 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6f\" d=\"M 1959 3097 \nQ 1497 3097 1228 2736 \nQ 959 2375 959 1747 \nQ 959 1119 1226 758 \nQ 1494 397 1959 397 \nQ 2419 397 2687 759 \nQ 2956 1122 2956 1747 \nQ 2956 2369 2687 2733 \nQ 2419 3097 1959 3097 \nz\nM 1959 3584 \nQ 2709 3584 3137 3096 \nQ 3566 2609 3566 1747 \nQ 3566 888 3137 398 \nQ 2709 -91 1959 -91 \nQ 1206 -91 779 398 \nQ 353 888 353 1747 \nQ 353 2609 779 3096 \nQ 1206 3584 1959 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-63\" d=\"M 3122 3366 \nL 3122 2828 \nQ 2878 2963 2633 3030 \nQ 2388 3097 2138 3097 \nQ 1578 3097 1268 2742 \nQ 959 2388 959 1747 \nQ 959 1106 1268 751 \nQ 1578 397 2138 397 \nQ 2388 397 2633 464 \nQ 2878 531 3122 666 \nL 3122 134 \nQ 2881 22 2623 -34 \nQ 2366 -91 2075 -91 \nQ 1284 -91 818 406 \nQ 353 903 353 1747 \nQ 353 2603 823 3093 \nQ 1294 3584 2113 3584 \nQ 2378 3584 2631 3529 \nQ 2884 3475 3122 3366 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-68\" d=\"M 3513 2113 \nL 3513 0 \nL 2938 0 \nL 2938 2094 \nQ 2938 2591 2744 2837 \nQ 2550 3084 2163 3084 \nQ 1697 3084 1428 2787 \nQ 1159 2491 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 4863 \nL 1159 4863 \nL 1159 2956 \nQ 1366 3272 1645 3428 \nQ 1925 3584 2291 3584 \nQ 2894 3584 3203 3211 \nQ 3513 2838 3513 2113 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-65\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(61.523438 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(186.181641 0)\"/>\n      <use xlink:href=\"#DejaVuSans-68\" transform=\"translate(241.162109 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_7\">\n      <defs>\n       <path id=\"m72c6ca78a9\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m72c6ca78a9\" x=\"30.103125\" y=\"142.201311\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- 0.0 -->\n      <g transform=\"translate(7.2 146.00053) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-2e\" d=\"M 684 794 \nL 1344 794 \nL 1344 0 \nL 684 0 \nL 684 794 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#m72c6ca78a9\" x=\"30.103125\" y=\"113.940124\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- 0.2 -->\n      <g transform=\"translate(7.2 117.739343) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_9\">\n      <g>\n       <use xlink:href=\"#m72c6ca78a9\" x=\"30.103125\" y=\"85.678937\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_10\">\n      <!-- 0.4 -->\n      <g transform=\"translate(7.2 89.478156) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-34\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_4\">\n     <g id=\"line2d_10\">\n      <g>\n       <use xlink:href=\"#m72c6ca78a9\" x=\"30.103125\" y=\"57.41775\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_11\">\n      <!-- 0.6 -->\n      <g transform=\"translate(7.2 61.216969) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-36\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_5\">\n     <g id=\"line2d_11\">\n      <g>\n       <use xlink:href=\"#m72c6ca78a9\" x=\"30.103125\" y=\"29.156563\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_12\">\n      <!-- 0.8 -->\n      <g transform=\"translate(7.2 32.955782) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-38\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_12\">\n    <path d=\"M 34.96476 57.581382 \nL 44.72976 91.733198 \nL 54.49476 103.272796 \nL 64.25976 104.969715 \nL 74.02476 112.382797 \nL 83.78976 113.03557 \nL 93.55476 119.055203 \nL 103.31976 118.721536 \nL 113.08476 124.537898 \nL 122.84976 123.377123 \nL 132.61476 129.066585 \nL 142.37976 128.14032 \nL 152.14476 132.556679 \nL 161.90976 132.672028 \nL 171.67476 135.397579 \nL 181.43976 135.06086 \nL 191.20476 138.089984 \nL 200.96976 137.638505 \nL 210.73476 139.5 \nL 220.49976 139.013261 \n\" clip-path=\"url(#p094447c306)\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n   </g>\n   <g id=\"line2d_13\">\n    <path d=\"M 49.633125 20.109742 \nL 69.163125 16.993521 \nL 88.693125 14.883434 \nL 108.223125 16.839806 \nL 127.753125 18.418877 \nL 147.283125 17.189158 \nL 166.813125 13.765508 \nL 186.343125 14.268575 \nL 205.873125 13.5 \nL 225.403125 16.127128 \n\" clip-path=\"url(#p094447c306)\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #ff7f0e; stroke-width: 1.5\"/>\n   </g>\n   <g id=\"line2d_14\">\n    <path d=\"M 49.633125 87.127665 \nL 69.163125 96.843451 \nL 88.693125 103.825718 \nL 108.223125 97.468608 \nL 127.753125 90.431294 \nL 147.283125 92.755802 \nL 166.813125 99.906076 \nL 186.343125 98.253573 \nL 205.873125 98.261354 \nL 225.403125 81.120725 \n\" clip-path=\"url(#p094447c306)\" style=\"fill: none; stroke-dasharray: 9.6,2.4,1.5,2.4; stroke-dashoffset: 0; stroke: #2ca02c; stroke-width: 1.5\"/>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 30.103125 145.8 \nL 30.103125 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 225.403125 145.8 \nL 225.403125 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 30.103125 145.8 \nL 225.403125 145.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 30.103125 7.2 \nL 225.403125 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"legend_1\">\n    <g id=\"patch_7\">\n     <path d=\"M 138.8125 60.06875 \nL 218.403125 60.06875 \nQ 220.403125 60.06875 220.403125 58.06875 \nL 220.403125 14.2 \nQ 220.403125 12.2 218.403125 12.2 \nL 138.8125 12.2 \nQ 136.8125 12.2 136.8125 14.2 \nL 136.8125 58.06875 \nQ 136.8125 60.06875 138.8125 60.06875 \nz\n\" style=\"fill: #ffffff; opacity: 0.8; stroke: #cccccc; stroke-linejoin: miter\"/>\n    </g>\n    <g id=\"line2d_15\">\n     <path d=\"M 140.8125 20.298438 \nL 150.8125 20.298438 \nL 160.8125 20.298438 \n\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n    </g>\n    <g id=\"text_13\">\n     <!-- train_loss -->\n     <g transform=\"translate(168.8125 23.798438) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-74\" d=\"M 1172 4494 \nL 1172 3500 \nL 2356 3500 \nL 2356 3053 \nL 1172 3053 \nL 1172 1153 \nQ 1172 725 1289 603 \nQ 1406 481 1766 481 \nL 2356 481 \nL 2356 0 \nL 1766 0 \nQ 1100 0 847 248 \nQ 594 497 594 1153 \nL 594 3053 \nL 172 3053 \nL 172 3500 \nL 594 3500 \nL 594 4494 \nL 1172 4494 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-72\" d=\"M 2631 2963 \nQ 2534 3019 2420 3045 \nQ 2306 3072 2169 3072 \nQ 1681 3072 1420 2755 \nQ 1159 2438 1159 1844 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1341 3275 1631 3429 \nQ 1922 3584 2338 3584 \nQ 2397 3584 2469 3576 \nQ 2541 3569 2628 3553 \nL 2631 2963 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-61\" d=\"M 2194 1759 \nQ 1497 1759 1228 1600 \nQ 959 1441 959 1056 \nQ 959 750 1161 570 \nQ 1363 391 1709 391 \nQ 2188 391 2477 730 \nQ 2766 1069 2766 1631 \nL 2766 1759 \nL 2194 1759 \nz\nM 3341 1997 \nL 3341 0 \nL 2766 0 \nL 2766 531 \nQ 2569 213 2275 61 \nQ 1981 -91 1556 -91 \nQ 1019 -91 701 211 \nQ 384 513 384 1019 \nQ 384 1609 779 1909 \nQ 1175 2209 1959 2209 \nL 2766 2209 \nL 2766 2266 \nQ 2766 2663 2505 2880 \nQ 2244 3097 1772 3097 \nQ 1472 3097 1187 3025 \nQ 903 2953 641 2809 \nL 641 3341 \nQ 956 3463 1253 3523 \nQ 1550 3584 1831 3584 \nQ 2591 3584 2966 3190 \nQ 3341 2797 3341 1997 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6e\" d=\"M 3513 2113 \nL 3513 0 \nL 2938 0 \nL 2938 2094 \nQ 2938 2591 2744 2837 \nQ 2550 3084 2163 3084 \nQ 1697 3084 1428 2787 \nQ 1159 2491 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1366 3272 1645 3428 \nQ 1925 3584 2291 3584 \nQ 2894 3584 3203 3211 \nQ 3513 2838 3513 2113 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-5f\" d=\"M 3263 -1063 \nL 3263 -1509 \nL -63 -1509 \nL -63 -1063 \nL 3263 -1063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6c\" d=\"M 603 4863 \nL 1178 4863 \nL 1178 0 \nL 603 0 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \nL 2834 2853 \nQ 2591 2978 2328 3040 \nQ 2066 3103 1784 3103 \nQ 1356 3103 1142 2972 \nQ 928 2841 928 2578 \nQ 928 2378 1081 2264 \nQ 1234 2150 1697 2047 \nL 1894 2003 \nQ 2506 1872 2764 1633 \nQ 3022 1394 3022 966 \nQ 3022 478 2636 193 \nQ 2250 -91 1575 -91 \nQ 1294 -91 989 -36 \nQ 684 19 347 128 \nL 347 722 \nQ 666 556 975 473 \nQ 1284 391 1588 391 \nQ 1994 391 2212 530 \nQ 2431 669 2431 922 \nQ 2431 1156 2273 1281 \nQ 2116 1406 1581 1522 \nL 1381 1569 \nQ 847 1681 609 1914 \nQ 372 2147 372 2553 \nQ 372 3047 722 3315 \nQ 1072 3584 1716 3584 \nQ 2034 3584 2315 3537 \nQ 2597 3491 2834 3397 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-74\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(39.208984 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(80.322266 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(141.601562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(169.384766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-5f\" transform=\"translate(232.763672 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(282.763672 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(310.546875 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(371.728516 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(423.828125 0)\"/>\n     </g>\n    </g>\n    <g id=\"line2d_16\">\n     <path d=\"M 140.8125 35.254688 \nL 150.8125 35.254688 \nL 160.8125 35.254688 \n\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #ff7f0e; stroke-width: 1.5\"/>\n    </g>\n    <g id=\"text_14\">\n     <!-- val_acc -->\n     <g transform=\"translate(168.8125 38.754688) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-76\" d=\"M 191 3500 \nL 800 3500 \nL 1894 563 \nL 2988 3500 \nL 3597 3500 \nL 2284 0 \nL 1503 0 \nL 191 3500 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-76\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(59.179688 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(120.458984 0)\"/>\n      <use xlink:href=\"#DejaVuSans-5f\" transform=\"translate(148.242188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(198.242188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(259.521484 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(314.501953 0)\"/>\n     </g>\n    </g>\n    <g id=\"line2d_17\">\n     <path d=\"M 140.8125 50.210938 \nL 150.8125 50.210938 \nL 160.8125 50.210938 \n\" style=\"fill: none; stroke-dasharray: 9.6,2.4,1.5,2.4; stroke-dashoffset: 0; stroke: #2ca02c; stroke-width: 1.5\"/>\n    </g>\n    <g id=\"text_15\">\n     <!-- val_loss -->\n     <g transform=\"translate(168.8125 53.710938) scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-76\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(59.179688 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(120.458984 0)\"/>\n      <use xlink:href=\"#DejaVuSans-5f\" transform=\"translate(148.242188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(198.242188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(226.025391 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(287.207031 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(339.306641 0)\"/>\n     </g>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"p094447c306\">\n   <rect x=\"30.103125\" y=\"7.2\" width=\"195.3\" height=\"138.6\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-10 fw=all -->\n",
    "\n",
    "## ResNeXt\n",
    "\n",
    "One of the challenges one encounters in the design of ResNet is the trade-off between nonlinearity and dimensionality within a given block. That is, we could add more nonlinearity by increasing the number of layers, or by increasing the width of the convolutions. An alternative strategy is to increase the number of channels that can carry information between blocks. Unfortunately, the latter comes with a quadratic penalty since the computational cost of ingesting $c_\\textrm{i}$ channels and emitting $c_\\textrm{o}$ channels is proportional to $\\mathcal{O}(c_\\textrm{i} \\cdot c_\\textrm{o})$ (see our discussion in that section). \n",
    "\n",
    "We can take some inspiration from the Inception block of the figure which has information flowing through the block in separate groups. Applying the idea of multiple independent groups to the ResNet block of the figure led to the design of ResNeXt [@Xie.Girshick.Dollar.ea.2017].\n",
    "Different from the smorgasbord of transformations in Inception, \n",
    "ResNeXt adopts the *same* transformation in all branches,\n",
    "thus minimizing the need for manual tuning of each branch. \n",
    "\n",
    "![The ResNeXt block: a grouped $3 \\times 3$ convolution with $\\mathit{g}$ groups, sandwiched between two $1 \\times 1$ convolutions that mix information across groups. Grouped convolution uses $\\mathit{g}$ times fewer multiply-adds than a dense convolution; the block is a bottleneck residual block when the number of intermediate channels $\\mathit{b}$ is less than $\\mathit{c}$.](../img/arch-resnext-block.svg)\n",
    "\n",
    "Breaking up a convolution from $c_\\textrm{i}$ to $c_\\textrm{o}$ channels into one of $g$ groups of size $c_\\textrm{i}/g$ generating $g$ outputs of size $c_\\textrm{o}/g$ is called, quite fittingly, a *grouped convolution*. The arithmetic cost is reduced from $\\mathcal{O}(c_\\textrm{i} \\cdot c_\\textrm{o})$ to $\\mathcal{O}(g \\cdot (c_\\textrm{i}/g) \\cdot (c_\\textrm{o}/g)) = \\mathcal{O}(c_\\textrm{i} \\cdot c_\\textrm{o} / g)$, i.e., it uses $g$ times fewer multiply-adds. The number of parameters is likewise reduced from a $c_\\textrm{i} \\times c_\\textrm{o}$ matrix to $g$ smaller matrices of size $(c_\\textrm{i}/g) \\times (c_\\textrm{o}/g)$. Wall-clock speedup depends on kernel efficiency and memory traffic and can be much smaller. In what follows we assume that both $c_\\textrm{i}$ and $c_\\textrm{o}$ are divisible by $g$.\n",
    "\n",
    "The only challenge in this design is that no information is exchanged between the $g$ groups. The ResNeXt block of \n",
    "the figure amends this in two ways: the grouped convolution with a $3 \\times 3$ kernel is sandwiched in between two $1 \\times 1$ convolutions. The second one serves double duty in changing the number of channels back. The benefit is that we only pay the $\\mathcal{O}(c \\cdot b)$ cost for $1 \\times 1$ kernels and can make do with an $\\mathcal{O}(b^2 / g)$ cost for $3 \\times 3$ kernels. Similar to the residual block implementation in\n",
    "that section, the residual connection is replaced (thus generalized) by a $1 \\times 1$ convolution.\n",
    "\n",
    "The block in the figure will also play a major role in the design of generic modern CNNs in that section. Grouped convolutions appeared in AlexNet [@Krizhevsky.Sutskever.Hinton.2012], where most channels were split into two groups so that each GPU could process one group with limited cross-GPU communication.\n",
    "\n",
    "The following implementation of the `ResNeXtBlock` class takes as argument `groups` ($g$), with \n",
    "`bot_channels` ($b$) intermediate (bottleneck) channels. Lastly, when we need to reduce the height and width of the representation, we add a stride of $2$ by setting `use_1x1conv=True, strides=2`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-resnext-1"
   },
   "source": [
    "class ResNeXtBlock(nnx.Module):\n",
    "    \"\"\"The ResNeXt block.\"\"\"\n",
    "    def __init__(self, num_channels, groups, bot_mul, use_1x1conv=False,\n",
    "                 strides=(1, 1), in_channels=None, rngs=None):\n",
    "        in_channels = num_channels if in_channels is None else in_channels\n",
    "        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs\n",
    "        bot_channels = int(round(num_channels * bot_mul))\n",
    "        self.conv1 = nnx.Conv(in_channels, bot_channels, kernel_size=(1, 1),\n",
    "                              strides=(1, 1), rngs=rngs)\n",
    "        self.conv2 = nnx.Conv(bot_channels, bot_channels,\n",
    "                              kernel_size=(3, 3), strides=strides,\n",
    "                              padding='same', feature_group_count=groups,\n",
    "                              rngs=rngs)\n",
    "        self.conv3 = nnx.Conv(bot_channels, num_channels,\n",
    "                              kernel_size=(1, 1), strides=(1, 1), rngs=rngs)\n",
    "        self.bn1 = nnx.BatchNorm(bot_channels, rngs=rngs)\n",
    "        self.bn2 = nnx.BatchNorm(bot_channels, rngs=rngs)\n",
    "        self.bn3 = nnx.BatchNorm(num_channels, rngs=rngs)\n",
    "        if use_1x1conv:\n",
    "            self.conv4 = nnx.Conv(in_channels, num_channels,\n",
    "                                  kernel_size=(1, 1), strides=strides,\n",
    "                                  rngs=rngs)\n",
    "            self.bn4 = nnx.BatchNorm(num_channels, rngs=rngs)\n",
    "        else:\n",
    "            self.conv4 = None\n",
    "\n",
    "    def __call__(self, X):\n",
    "        Y = nnx.relu(self.bn1(self.conv1(X)))\n",
    "        Y = nnx.relu(self.bn2(self.conv2(Y)))\n",
    "        Y = self.bn3(self.conv3(Y))\n",
    "        if self.conv4:\n",
    "            X = self.bn4(self.conv4(X))\n",
    "        return nnx.relu(Y + X)"
   ],
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-11 fw=all -->\n",
    "\n",
    "Its use is entirely analogous to that of the `ResNetBlock` discussed previously. For instance, when using (`use_1x1conv=False, strides=1`), the input and output are of the same shape. Alternatively, setting `use_1x1conv=True, strides=2` halves the output height and width."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "resnet-resnext-2"
   },
   "source": [
    "blk = ResNeXtBlock(32, 16, 1)\n",
    "X = jnp.zeros((4, 96, 96, 32))\n",
    "blk(X).shape"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(4, 96, 96, 32)"
     },
     "metadata": {},
     "execution_count": 12
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-12 fw=all -->\n",
    "\n",
    "## Concatenation instead of Addition: DenseNet\n",
    "\n",
    "The residual block computes $f(\\mathbf{x}) = \\mathbf{x} + g(\\mathbf{x})$, splitting a function into the identity plus a correction, much as a Taylor expansion splits a function into a leading term plus higher-order refinements. What if we want to keep more than two terms, without forcing them to share a sum? *DenseNet* (dense convolutional network) [@Huang.Liu.Van-Der-Maaten.ea.2017] answers by *concatenating* a layer's output onto its input along the channel dimension, so every layer receives the feature maps of all layers that precede it. Applying an increasingly complex sequence of functions thus maps $\\mathbf{x}$ to\n",
    "\n",
    "$$\n",
    "\\mathbf{x} \\to \\left[\n",
    "\\mathbf{x},\n",
    "f_1(\\mathbf{x}),\n",
    "f_2\\left(\\left[\\mathbf{x}, f_1\\left(\\mathbf{x}\\right)\\right]\\right), f_3\\left(\\left[\\mathbf{x}, f_1\\left(\\mathbf{x}\\right), f_2\\left(\\left[\\mathbf{x}, f_1\\left(\\mathbf{x}\\right)\\right]\\right)\\right]\\right), \\ldots\\right],\n",
    "$$\n",
    "\n",
    "where $[\\cdot, \\cdot]$ denotes concatenation. The name reflects the resulting dependency graph: the last layer of such a chain is densely connected to all its predecessors. the figure contrasts the two designs.\n",
    "\n",
    "![A residual block (left) adds its input back to the output, keeping the number of channels fixed. A DenseNet block (right) concatenates each convolution's output onto its input, so the number of channels grows with every layer while all earlier features remain directly accessible.](../img/arch-densenet-block.svg)\n",
    "\n",
    "A DenseNet consists of *dense blocks*, which define how outputs are concatenated, alternating with *transition layers*, which shrink the accumulated channels back down. Each convolution inside a dense block uses the \"batch normalization, activation, and convolution\" pre-activation ordering explored in the exercises below."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "densenet-dense-blocks-1"
   },
   "source": [
    "class ConvBlock(nnx.Module):\n",
    "    def __init__(self, in_channels, num_channels, rngs):\n",
    "        self.bn = nnx.BatchNorm(in_channels, rngs=rngs)\n",
    "        self.conv = nnx.Conv(in_channels, num_channels, kernel_size=(3, 3),\n",
    "                             padding=(1, 1), rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        return self.conv(nnx.relu(self.bn(X)))"
   ],
   "execution_count": 13,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-13 fw=all -->\n",
    "\n",
    "A *dense block* consists of multiple convolution blocks, each using the same number of output channels. In the forward propagation, however, we concatenate the input and output of each convolution block on the channel dimension. Lazy evaluation allows us to adjust the dimensionality automatically."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "densenet-dense-blocks-2"
   },
   "source": [
    "class DenseBlock(nnx.Module):\n",
    "    def __init__(self, num_convs, num_channels, in_channels=3, rngs=None):\n",
    "        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs\n",
    "        layers = []\n",
    "        for i in range(num_convs):\n",
    "            layers.append(ConvBlock(\n",
    "                in_channels + i * num_channels, num_channels, rngs))\n",
    "        self.layers = nnx.List(layers)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        for layer in self.layers:\n",
    "            Y = layer(X)\n",
    "            # Concatenate input and output of each block along the channels\n",
    "            X = jnp.concatenate((X, Y), axis=-1)\n",
    "        return X"
   ],
   "execution_count": 14,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-14 fw=all -->\n",
    "\n",
    "In the following example,\n",
    "we define a `DenseBlock` instance with two convolution blocks of 10 output channels.\n",
    "When using an input with three channels, we will get an output with  $3 + 10 + 10=23$ channels. The number of convolution block channels controls the growth in the number of output channels relative to the number of input channels. This is also referred to as the *growth rate*."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "densenet-dense-blocks-3"
   },
   "source": [
    "blk = DenseBlock(2, 10)\n",
    "X = jnp.zeros((4, 8, 8, 3))\n",
    "Y = blk(X)\n",
    "Y.shape"
   ],
   "execution_count": 15,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(4, 8, 8, 23)"
     },
     "metadata": {},
     "execution_count": 15
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-15 fw=all -->\n",
    "\n",
    "Since each dense block increases the number of channels, stacking many of them without correction would produce an excessively wide model. A *transition layer* reduces the number of channels by a $1\\times 1$ convolution and halves the height and width via average pooling with a stride of 2."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "densenet-transition-layers-1"
   },
   "source": [
    "class TransitionBlock(nnx.Module):\n",
    "    def __init__(self, in_channels, num_channels, rngs=None):\n",
    "        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs\n",
    "        self.bn = nnx.BatchNorm(in_channels, rngs=rngs)\n",
    "        self.conv = nnx.Conv(in_channels, num_channels,\n",
    "                             kernel_size=(1, 1), rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        X = self.conv(nnx.relu(self.bn(X)))\n",
    "        X = nnx.avg_pool(X, window_shape=(2, 2), strides=(2, 2))\n",
    "        return X"
   ],
   "execution_count": 16,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-16 fw=all -->\n",
    "\n",
    "Apply a transition layer with 10 channels to the output of the dense block in the previous example.  This reduces the number of output channels to 10, and halves the height and width."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "densenet-transition-layers-2"
   },
   "source": [
    "blk = TransitionBlock(23, 10)\n",
    "blk(Y).shape"
   ],
   "execution_count": 17,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(4, 4, 4, 10)"
     },
     "metadata": {},
     "execution_count": 17
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-17 fw=all -->\n",
    "\n",
    "A full DenseNet alternates four dense blocks with transition layers, mirroring\n",
    "the four-stage layout of ResNet-18; assembling and training one is a\n",
    "straightforward variation on the ResNet code above. Feature reuse makes\n",
    "DenseNet parameter-efficient: it reached ResNet-level ImageNet accuracy with\n",
    "fewer parameters. Concatenation carries a cost that addition does not. The\n",
    "unique feature maps grow linearly with depth, but a naive implementation\n",
    "repeatedly materializes ever-wider concatenations and temporary normalized\n",
    "outputs, producing quadratic memory growth inside a dense block.\n",
    "Memory-efficient implementations recompute those temporaries during\n",
    "backpropagation [@pleiss2017memory]. The remaining activation traffic\n",
    "helped residual addition become the more common choice at scale.\n",
    "\n",
    "## Summary and Discussion\n",
    "\n",
    "Nested function classes are desirable because adding capacity then preserves\n",
    "the functions already available. A residual branch parametrizes a block as the\n",
    "identity plus a correction. In a pre-activation block, setting that correction\n",
    "to zero gives the identity exactly; in the post-activation block used above,\n",
    "the statement holds on the nonnegative inputs produced by the preceding ReLU.\n",
    "\n",
    "\n",
    "The residual mapping makes a near-identity transformation accessible by\n",
    "driving the correction toward zero, while the shortcut supplies a direct path\n",
    "for activations and gradients. This enabled the original ResNet experiments\n",
    "to train networks as deep as 152 layers [@He.Zhang.Ren.ea.2016].\n",
    "Residual branches can also be initialized near the identity, for example by\n",
    "zero-initializing the final normalization scale. Ordinary randomly initialized\n",
    "residual blocks are not exact identities, and inserting new blocks during\n",
    "training requires a separate function-preserving procedure.\n",
    "\n",
    "Prior to residual connections,\n",
    "bypassing paths with gating units were introduced\n",
    "to effectively train highway networks with over 100 layers\n",
    "[@srivastava2015highway].\n",
    "Using identity functions as bypassing paths,\n",
    "ResNet performed well\n",
    "on multiple computer vision tasks.\n",
    "Residual connections had a major influence on the design of subsequent deep neural networks, of either convolutional or sequential nature.\n",
    "As we will introduce later,\n",
    "the Transformer architecture [@Vaswani.Shazeer.Parmar.ea.2017]\n",
    "adopts residual connections (together with other design choices) and is pervasive\n",
    "in areas as diverse as\n",
    "language, vision, speech, and reinforcement learning.\n",
    "Every transformer block in a large language model, and every block of the\n",
    "denoising networks behind diffusion models, is a residual block. By parameter\n",
    "count, most residual blocks in the world now live outside convolutional\n",
    "networks.\n",
    "\n",
    "ResNeXt illustrates the trade between arithmetic and activation width. A\n",
    "grouped convolution is a block-diagonal channel-mixing matrix, reducing\n",
    "multiply-adds while retaining a wide representation. ShiftNet\n",
    "[@wu2018shift] takes this idea further by replacing spatial\n",
    "multiplications with shifts of channel activations. The shifts use no\n",
    "floating-point multiply-adds, though moving the activation data still incurs\n",
    "memory traffic.\n",
    "\n",
    "A common feature of the designs we have discussed so far is that the network design is fairly manual, primarily relying on the ingenuity of the designer to find the \"right\" network hyperparameters. While clearly feasible, it is also very costly in terms of human time and there is no guarantee that the outcome is optimal in any sense. In that section we will discuss a number of strategies for obtaining high quality networks in a more automated fashion. In particular, we will review the notion of *network design spaces* that led to the RegNetX/Y models\n",
    "[@Radosavovic.Kosaraju.Girshick.ea.2020].\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. What are the major differences between the Inception block in the figure and the residual block? How do they compare in terms of computation, accuracy, and the classes of functions they can describe?\n",
    "1. Refer to Table 1 in the ResNet paper [@He.Zhang.Ren.ea.2016] to implement different variants of the network. \n",
    "1. For deeper networks, ResNet introduces a \"bottleneck\" architecture to reduce model complexity. Try to implement it.\n",
    "1. In subsequent versions of ResNet, the authors changed the \"convolution, batch normalization, and activation\" structure to the \"batch normalization, activation, and convolution\" structure. Make this improvement yourself. See Figure 1 in @He.Zhang.Ren.ea.2016*1 for details. This ordering is essentially what ConvNeXt adopts (that section).\n",
    "1. Why can't we just increase the complexity of functions without bound, even if the function classes are nested?\n",
    "1. One of the advantages claimed in the DenseNet paper [@Huang.Liu.Van-Der-Maaten.ea.2017] is that its models have fewer parameters than comparable ResNets. Why is this the case? Consider which computations a concatenated feature saves relative to recomputing it.\n",
    "1. For a dense block whose $k$ convolutions each emit $g$ channels (the growth rate) on an input with $c$ channels, how many channels does the $i$-th convolution consume? Sum these to compare the activation memory of the dense block with that of $k$ residual blocks of constant width $c$, and relate your answer to the memory-efficient implementations of @pleiss2017memory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=resnet-md-18 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/18006)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
