{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-1 fw=all -->\n",
    "\n",
    "# Modules and Model Construction\n",
    "\n",
    "The networks we have trained so far were small enough to write down layer by\n",
    "layer. The networks ahead are not: deep ResNets chain more than a hundred\n",
    "convolutional layers [@He.Zhang.Ren.ea.2016], and a GPT-style language\n",
    "model stacks dozens of identical Transformer blocks\n",
    "[@Radford.Wu.Child.ea.2019]. Such models are assembled from repeated\n",
    "*blocks*, groups of layers treated as units of design. The abstraction that\n",
    "makes blocks composable is the *module*.\n",
    "\n",
    "A module is an object with three responsibilities: it owns *parameters*, it\n",
    "owns *child modules*, and it implements a *forward computation* that maps\n",
    "inputs to outputs. The definition is deliberately recursive. A fully connected\n",
    "layer is a module (parameters, no children). A residual block is a module (no\n",
    "direct parameters, a few child layers that own parameters). A hundred-layer\n",
    "network is a module whose children are blocks whose children are layers. Most\n",
    "models therefore have a tree-shaped module hierarchy, as sketched in\n",
    "the figure. Reusing one child at several sites turns that tree into\n",
    "an object graph, a case we meet when tying parameters in\n",
    "that section. Almost everything\n",
    "this chapter does to a model, listing its parameters\n",
    "(that section), moving it to a GPU (that section),\n",
    "saving it to disk (that section), is implemented as a walk over\n",
    "that tree.\n",
    "\n",
    "![Layers compose into blocks and blocks compose into models, giving the usual tree-shaped module hierarchy.](../img/bg-module-tree.svg)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-modules-and-model-construction"
   },
   "source": [
    "from dataclasses import dataclass\n",
    "from mxnet import init, np, npx\n",
    "from mxnet.gluon import nn\n",
    "npx.set_np()"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-2 fw=all -->\n",
    "\n",
    "## The Module Abstraction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-3 fw=mxnet -->\n",
    "\n",
    "In MXNet the module class is Gluon's `nn.Block`. We have used one of its\n",
    "subclasses all along: `nn.Sequential` builds a model from a chain of layers\n",
    "appended with its `add` method, here the familiar MLP with a 256-unit ReLU\n",
    "hidden layer and a 10-unit output layer. Gluon attaches the activation to the\n",
    "`Dense` layer itself rather than interposing a separate activation layer, and\n",
    "one `initialize()` call stands between construction and the first forward\n",
    "pass; what that call does, and what it defers, is where this section ends."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-1"
   },
   "source": [
    "net = nn.Sequential()\n",
    "net.add(nn.Dense(256, activation='relu'))\n",
    "net.add(nn.Dense(10))\n",
    "net.initialize()\n",
    "\n",
    "X = np.random.uniform(size=(2, 20))\n",
    "net(X).shape"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 10)"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-4 fw=mxnet -->\n",
    "\n",
    "`Sequential` is not a special construct. It is itself a `Block` whose forward\n",
    "computation runs its children in order, and the children are the two layers\n",
    "`add` appended, stored under names in a registry that printing the model\n",
    "displays:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-2"
   },
   "source": [
    "net"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Sequential(\n  (0): Dense(20 -> 256, Activation(relu))\n  (1): Dense(256 -> 10, linear)\n)"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-5 fw=mxnet -->\n",
    "\n",
    "This registry is what the `Block` machinery traverses: `net.collect_params()`\n",
    "gathers parameters by walking the children recursively and keys them by those\n",
    "registry names (`0.weight`, `1.bias`, ...), and the same walk underlies\n",
    "initialization, device movement, and serialization. A block the registry does\n",
    "not contain might as well not exist, a fact we will exploit for a\n",
    "demonstration shortly."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-3"
   },
   "source": [
    "class MLP(nn.Block):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.hidden = nn.Dense(256, activation='relu')\n",
    "        self.out = nn.Dense(10)\n",
    "\n",
    "    def forward(self, X):\n",
    "        return self.out(self.hidden(X))"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-4"
   },
   "source": [
    "net = MLP()\n",
    "net.initialize()\n",
    "net(X).shape"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 10)"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-6 fw=mxnet -->\n",
    "\n",
    "Two details do the work here. First, `self.hidden = nn.Dense(...)` is not an\n",
    "ordinary attribute assignment: `nn.Block` intercepts `__setattr__`, sees that\n",
    "the value is a block, and adds it to the registry we just inspected. That is\n",
    "why both layers' parameters show up in `net.collect_params()` with no further\n",
    "ceremony. Second, we never wrote a backward method; automatic differentiation\n",
    "derives gradients from whatever `forward` computes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-7 fw=all -->\n",
    "\n",
    "## Sequential and Friends: Containers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-8 fw=mxnet -->\n",
    "\n",
    "To see that there is no magic left in `nnx.Sequential`, we can write it\n",
    "ourselves. Two ingredients suffice: register each child, and loop over the\n",
    "children in `forward`. One Gluon 2.0 wrinkle: the registry holds weak\n",
    "references, so our class also keeps the blocks in a plain list to keep them\n",
    "alive."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-1"
   },
   "source": [
    "class MySequential(nn.Block):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # The registry holds weakrefs; this list keeps the blocks alive\n",
    "        self._layers = []\n",
    "\n",
    "    def add(self, block):\n",
    "        self._layers.append(block)\n",
    "        self.register_child(block)\n",
    "\n",
    "    def forward(self, X):\n",
    "        for block in self._layers:\n",
    "            X = block(X)\n",
    "        return X"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-9 fw=mxnet -->\n",
    "\n",
    "`register_child` writes a child into the registry under a string name (that\n",
    "is where the `(0)` and `(1)` keys above came from; attribute assignment calls\n",
    "it with the attribute name instead). Our version is a drop-in replacement:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-2"
   },
   "source": [
    "net = MySequential()\n",
    "net.add(nn.Dense(256, activation='relu'))\n",
    "net.add(nn.Dense(10))\n",
    "net.initialize()\n",
    "net(X).shape"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 10)"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-10 fw=mxnet -->\n",
    "\n",
    "The registration step is easy to lose. The following block stores its layers\n",
    "in a plain Python list, which the `__setattr__` interception ignores, so\n",
    "nothing inside it is registered. Since `net.initialize()` walks the registry,\n",
    "it cannot reach the hidden children either, so we initialize each layer by\n",
    "hand just to get a forward pass, and the model appears to work:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-3"
   },
   "source": [
    "class PlainListMLP(nn.Block):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.layers = [nn.Dense(256, activation='relu'), nn.Dense(10)]\n",
    "        for layer in self.layers:\n",
    "            layer.initialize()  # By hand: net.initialize() cannot see them\n",
    "\n",
    "    def forward(self, X):\n",
    "        for layer in self.layers:\n",
    "            X = layer(X)\n",
    "        return X\n",
    "\n",
    "net = PlainListMLP()\n",
    "net(X).shape, len(net.collect_params())"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/tmp/ipykernel_65430/1272845916.py:14: UserWarning: \"PlainListMLP.layers\" is an unregistered container with Blocks. Note that Blocks inside the list, tuple or dict will not be registered automatically. Make sure to register them using register_child() or switching to nn.Sequential/nn.HybridSequential instead. \n  net(X).shape, len(net.collect_params())\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "((2, 10), 0)"
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-11 fw=mxnet -->\n",
    "\n",
    "The model computes, yet it owns zero parameters: hand it to a `Trainer` and\n",
    "training updates nothing; `save_parameters` checkpoints an empty model. Gluon\n",
    "does not let this pass in silence, though. `collect_params()` scans the\n",
    "block's ordinary attributes for lists, tuples, and dictionaries holding\n",
    "unregistered blocks, and the call above printed a warning naming the guilty\n",
    "attribute (`\"PlainListMLP.layers\" is an unregistered container with\n",
    "Blocks...`) along with the fix. And had we skipped the by-hand\n",
    "initialization, the first forward pass would have stopped with a\n",
    "`Parameter has not been initialized` error instead of running at all. Where\n",
    "PyTorch's version of this bug is diagnosed by staring at a loss curve,\n",
    "Gluon's announces itself."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-4"
   },
   "source": [
    "class RegisteredListMLP(PlainListMLP):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        for layer in self.layers:\n",
    "            self.register_child(layer)\n",
    "\n",
    "net = RegisteredListMLP()\n",
    "net(X).shape, sum(p.data().size for p in net.collect_params().values())"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "((2, 10), 7946)"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-12 fw=mxnet -->\n",
    "\n",
    "Same forward pass, 7946 registered parameters, and no warning. Gluon ships no\n",
    "`ModuleList` or `ModuleDict` equivalents: the two registered homes for\n",
    "children are named attributes and `nn.Sequential`, with `register_child` as\n",
    "the escape hatch for anything else, such as a list of blocks whose loop body\n",
    "is not a plain chain. Transformer implementations in Gluon conventionally\n",
    "keep their stack of blocks in an `nn.Sequential` (or register them one by\n",
    "one) and their named parts (embedding, final normalization, output head) as\n",
    "attributes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-13 fw=all -->\n",
    "\n",
    "## Forward Is Just Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-14 fw=mxnet -->\n",
    "\n",
    "`forward` is an ordinary Python method. Nothing restricts it to chaining\n",
    "children: it can branch, loop, call any `np` function, and combine\n",
    "intermediate results however it likes; a Gluon `Block` executes eagerly, so\n",
    "all of this runs one operation at a time, just as in NumPy. (Gluon's\n",
    "`HybridBlock` can compile the forward computation into a graph for speed, at\n",
    "the price of restricting it to traceable operations; we stay with `Block`\n",
    "here.) The loop in `MySequential` already used this freedom. Its most\n",
    "consequential one-line use is the *residual connection*, the wiring idiom at\n",
    "the heart of ResNets and Transformers alike:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-forward-is-just-python-1"
   },
   "source": [
    "class ResidualBlock(nn.Block):\n",
    "    def __init__(self, num_hiddens):\n",
    "        super().__init__()\n",
    "        self.body = nn.Sequential()\n",
    "        self.body.add(nn.Dense(num_hiddens, activation='relu'),\n",
    "                      nn.Dense(num_hiddens))\n",
    "\n",
    "    def forward(self, X):\n",
    "        return X + self.body(X)"
   ],
   "execution_count": 10,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-15 fw=all -->\n",
    "\n",
    "![The residual wiring `X + body(X)`: the input splits at a branch point into the body stack and an identity skip, and the two rejoin by addition before the block's output.](../img/bg-residual-block.svg)\n",
    "\n",
    "`X + body(X)` is not a layer any framework provides. It is plain arithmetic\n",
    "in the forward computation, and it changes what the block *is*: the block\n",
    "computes a perturbation of the identity function rather than an arbitrary\n",
    "transformation. If its body is $F$, its Jacobian is $I + J_F$, so\n",
    "backpropagation receives an additive identity contribution along the skip\n",
    "path. This contribution can still cancel against $J_F$; it is a direct path,\n",
    "not a guarantee that every gradient is preserved\n",
    "[@He.Zhang.Ren.ea.2016]. the figure diagrams\n",
    "exactly this wiring, and that section develops both points when we\n",
    "build ResNet; for now we only need the mechanics. One mechanical consequence\n",
    "is visible already: the addition forces the input and output shapes to\n",
    "agree, so a residual block has a single width that is part of its identity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-16 fw=mxnet -->\n",
    "\n",
    "Gluon defers input widths to the first forward pass (a mechanism we examine\n",
    "at the end of this section), so here the constraint falls on the caller:\n",
    "feed the block anything other than `num_hiddens` columns and the addition\n",
    "fails."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-forward-is-just-python-2"
   },
   "source": [
    "block = ResidualBlock(24)\n",
    "block.initialize()\n",
    "block(np.random.normal(size=(2, 24))).shape"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 24)"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-17 fw=mxnet -->\n",
    "\n",
    "`forward` may also use state that is neither an input nor a parameter.\n",
    "Suppose we want to damp each block's contribution by a fixed factor:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-forward-is-just-python-3"
   },
   "source": [
    "class ScaledResidual(ResidualBlock):\n",
    "    def __init__(self, num_hiddens, alpha=0.5):\n",
    "        super().__init__(num_hiddens)\n",
    "        self.alpha = np.array(alpha)  # Fixed by design, never trained\n",
    "\n",
    "    def forward(self, X):\n",
    "        return X + self.alpha * self.body(X)\n",
    "\n",
    "block = ScaledResidual(24)\n",
    "'alpha' in block.collect_params(), list(block.collect_params())[:2]"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(False, ['body.0.weight', 'body.0.bias'])"
     },
     "metadata": {},
     "execution_count": 12
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-18 fw=mxnet -->\n",
    "\n",
    "`alpha` enters the computation, but it is not a parameter: as the output\n",
    "shows, it never appears in `collect_params()`, so the `Trainer` never touches\n",
    "it. That much we wanted. Storing it as a plain attribute has a cost we did\n",
    "not want, though: invisible to the parameter walk, it will not be saved by\n",
    "`save_parameters`, and it will not move to a GPU with the rest of the block.\n",
    "Some state is not a parameter but must still travel with the model; the\n",
    "registered home for such state is a constant parameter (`gluon.Constant`),\n",
    "introduced in that section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-19 fw=all -->\n",
    "\n",
    "## Lazy Initialization: Shapes from Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-20 fw=mxnet -->\n",
    "\n",
    "We have been doing something odd since our first MLP without commenting on\n",
    "it: `nn.Dense(256)` names only the layer's *output* width. Its weight matrix\n",
    "has shape `(256, in_units)`, and we never said what `in_units` is. The layer\n",
    "cannot know it at construction time, since it depends on the data it will\n",
    "receive. So Gluon does not allocate parameters at construction time at all.\n",
    "Deferred initialization is not a special mode or a dedicated `Lazy*` class;\n",
    "it is how every Gluon layer works, and `in_units` is simply an optional\n",
    "argument:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-lazy-initialization-shapes-from-data-1"
   },
   "source": [
    "net = nn.Sequential()\n",
    "net.add(nn.Dense(256, activation='relu'))\n",
    "net.add(nn.Dense(10))\n",
    "net.collect_params()"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'0.weight': Parameter (shape=(256, -1), dtype=float32),\n '0.bias': Parameter (shape=(256,), dtype=float32),\n '1.weight': Parameter (shape=(10, -1), dtype=float32),\n '1.bias': Parameter (shape=(10,), dtype=float32)}"
     },
     "metadata": {},
     "execution_count": 13
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-21 fw=mxnet -->\n",
    "\n",
    "Each weight is a placeholder: the parameter objects exist, but the special\n",
    "value -1 marks the input dimension as unknown, and reading\n",
    "`net[0].weight.data()` would raise an error telling us to initialize the\n",
    "network first. Even `initialize()` allocates nothing at this point; it only\n",
    "records which initializer (and device) to use once the shapes are known. The\n",
    "first time data flows through, each layer reads its input width from the\n",
    "incoming shape, allocates and initializes a real weight, and its declared\n",
    "output width fixes the input of the next layer, so shapes cascade through the\n",
    "whole model:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-lazy-initialization-shapes-from-data-2"
   },
   "source": [
    "net.initialize()\n",
    "net(X)\n",
    "net.collect_params()"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'0.weight': Parameter (shape=(256, 20), dtype=float32),\n '0.bias': Parameter (shape=(256,), dtype=float32),\n '1.weight': Parameter (shape=(10, 256), dtype=float32),\n '1.bias': Parameter (shape=(10,), dtype=float32)}"
     },
     "metadata": {},
     "execution_count": 14
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-22 fw=all -->\n",
    "\n",
    "Lazy layers can remove shape arithmetic from model definitions. In a\n",
    "convolutional network the flattened feature-map size depends on the input\n",
    "resolution and upstream strides. We usually avoid that fragile arithmetic by\n",
    "using global or adaptive pooling before the first linear layer. When a width\n",
    "is part of the architecture, we name it in the model config and pass it to\n",
    "adjacent layers — code that stays explicit and compact without hiding\n",
    "parameter creation behind a sample batch, whether or not the framework\n",
    "offers a lazy mode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-23 fw=mxnet -->\n",
    "\n",
    "Until the first forward pass, Gluon parameters have placeholder shapes and no\n",
    "allocated values. A `Trainer` can be constructed over those parameter\n",
    "objects, but reading a weight or counting elements requires a *dry run* on a\n",
    "representative batch. A related subtlety: the random\n",
    "draw now happens at first call rather than at `initialize()`, so any random\n",
    "numbers your program draws in between shift the generator's state, and a\n",
    "fixed seed can yield different weights than an explicitly shaped version of\n",
    "the same model (that section returns to seeding). Other libraries\n",
    "retrofitted this behavior; Gluon was designed around it, which is why\n",
    "`initialize()` is a separate, explicit step rather than something the\n",
    "constructor does."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-lazy-initialization-shapes-from-data-4"
   },
   "source": [
    "net = nn.Sequential()\n",
    "net.add(nn.Dense(256, activation='relu'))\n",
    "net.add(nn.Dense(10))\n",
    "net.initialize(init.Xavier())\n",
    "net(X)\n",
    "net[0].weight.data().shape"
   ],
   "execution_count": 15,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(256, 20)"
     },
     "metadata": {},
     "execution_count": 15
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-24 fw=mxnet -->\n",
    "\n",
    "The dry run materialized every shape and drew the first weight's initial\n",
    "values from the Xavier initializer that `initialize` had recorded. Which\n",
    "initializer to use, and why Xavier's variance rule (that section)\n",
    "is a sensible default, is the subject of that section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-25 fw=all -->\n",
    "\n",
    "## Building from a Config\n",
    "\n",
    "So far every width in this section was a literal typed into a constructor.\n",
    "Real model code does not work that way. The architecture choices that vary\n",
    "across experiments, such as depth, width, and output size, must be logged with\n",
    "results and stored with checkpoints so that a saved model can be rebuilt. The\n",
    "topology remains in code; a small configuration object records its variable\n",
    "choices:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-1"
   },
   "source": [
    "@dataclass\n",
    "class MLPConfig:\n",
    "    d_in: int = 784\n",
    "    d_hidden: int = 256\n",
    "    num_blocks: int = 4\n",
    "    d_out: int = 10\n",
    "\n",
    "def build(cfg: MLPConfig) -> nn.Block:\n",
    "    net = nn.Sequential()\n",
    "    net.add(nn.Dense(cfg.d_hidden, in_units=cfg.d_in))\n",
    "    for _ in range(cfg.num_blocks):\n",
    "        net.add(ResidualBlock(cfg.d_hidden))\n",
    "    net.add(nn.Dense(cfg.d_out, in_units=cfg.d_hidden))\n",
    "    return net"
   ],
   "execution_count": 16,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-26 fw=mxnet -->\n",
    "\n",
    "One config produces one architecture: an input projection into the hidden\n",
    "width, `num_blocks` identical residual blocks, and an output projection. This\n",
    "is where explicit shapes beat deferred ones: the config already knows every\n",
    "width, so the two projections take `in_units` directly and skip deferred\n",
    "initialization. (The `Dense` layers inside each `ResidualBlock` were written\n",
    "without `in_units`, so the model as a whole still wants one dry run before\n",
    "its parameters can be counted.) Printing the model displays the module tree:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-2"
   },
   "source": [
    "net = build(MLPConfig())\n",
    "net.initialize()\n",
    "net"
   ],
   "execution_count": 17,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Sequential(\n  (0): Dense(784 -> 256, linear)\n  (1): ResidualBlock(\n    (body): Sequential(\n      (0): Dense(-1 -> 256, Activation(relu))\n      (1): Dense(-1 -> 256, linear)\n    )\n  )\n  (2): ResidualBlock(\n    (body): Sequential(\n      (0): Dense(-1 -> 256, Activation(relu))\n      (1): Dense(-1 -> 256, linear)\n    )\n  )\n  (3): ResidualBlock(\n    (body): Sequential(\n      (0): Dense(-1 -> 256, Activation(relu))\n      (1): Dense(-1 -> 256, linear)\n    )\n  )\n  (4): ResidualBlock(\n    (body): Sequential(\n      (0): Dense(-1 -> 256, Activation(relu))\n      (1): Dense(-1 -> 256, linear)\n    )\n  )\n  (5): Dense(256 -> 10, linear)\n)"
     },
     "metadata": {},
     "execution_count": 17
    }
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-3"
   },
   "source": [
    "net(np.random.uniform(size=(2, 784))).shape"
   ],
   "execution_count": 18,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 10)"
     },
     "metadata": {},
     "execution_count": 18
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-27 fw=all -->\n",
    "\n",
    "The variable architecture choices are now data. Rescaling the model is a change to two fields, not\n",
    "an edit to model code:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-4"
   },
   "source": [
    "for cfg in (MLPConfig(), MLPConfig(d_hidden=512, num_blocks=8)):\n",
    "    net = build(cfg)\n",
    "    net.initialize()\n",
    "    net(np.zeros((2, cfg.d_in)))  # Dry run: materialize the deferred shapes\n",
    "    n = sum(p.data().size for p in net.collect_params().values())\n",
    "    print(f'd_hidden={cfg.d_hidden}, num_blocks={cfg.num_blocks}: '\n",
    "          f'{n:,} parameters')"
   ],
   "execution_count": 19,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "d_hidden=256, num_blocks=4: 729,866 parameters\nd_hidden=512, num_blocks=8: 4,609,546 parameters\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-28 fw=mxnet -->\n",
    "\n",
    "Because `build` is deterministic in `cfg`, the config is all you need to\n",
    "reconstruct the module tree later; that section saves it\n",
    "alongside the weights so that loading a checkpoint starts by rebuilding the\n",
    "exact same model. A config of widths and depths feeding a loop that stacks\n",
    "identical residual blocks is a common construction pattern in the Transformer\n",
    "implementations later in the book."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-29 fw=all -->\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-30 fw=mxnet -->\n",
    "\n",
    "A module owns parameters, child blocks, and a `forward` method. Layers,\n",
    "blocks, and whole models are the same kind of object, so the usual hierarchy\n",
    "is tree-shaped, while shared children introduce aliases. Parameter collection,\n",
    "initialization, device movement, and serialization traverse that object graph.\n",
    "Children are discovered through\n",
    "registration: attribute assignment, `nn.Sequential`'s `add`, and\n",
    "`register_child` register children, while a plain Python list hides them,\n",
    "though `collect_params()` warns about the hidden blocks by name. `forward` is\n",
    "ordinary Python; a residual connection is one line in it. Deferred\n",
    "initialization is Gluon's native mode: every layer declares its output width,\n",
    "`initialize()` records the initializer, and the first forward pass infers\n",
    "input widths and allocates the parameters, so inspection and training follow\n",
    "a dry run. Configs turn architecture into data: a `dataclass` of widths and\n",
    "depths plus a `build` function that stacks blocks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-31 fw=all -->\n",
    "\n",
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-32 fw=mxnet -->\n",
    "\n",
    "1. Take `PlainListMLP` and catalog everything that breaks besides the empty\n",
    "   parameter dictionary. Check `net.save_parameters(...)`, what\n",
    "   `net.initialize()` reaches, and what happens on the first forward pass if\n",
    "   you remove the by-hand initialization from the constructor. Explain how\n",
    "   each failure follows from the same missing registration, and note where\n",
    "   along the way the `collect_params()` warning fires.\n",
    "1. Implement a `ParallelBlock` that takes two child blocks `net1` and `net2`,\n",
    "   runs both on the same input, and concatenates their outputs along the last\n",
    "   dimension. What must be true of the two children's outputs for the\n",
    "   concatenation to be valid?\n",
    "1. Extend `MLPConfig` with an activation switch (for example,\n",
    "   `act: str = 'relu'`) and make `build` honor it. Which decisions belong in a\n",
    "   config and which belong in code? Where would you put a choice between\n",
    "   `ResidualBlock` and a plain feed-forward block?\n",
    "1. `ResidualBlock` requires its input and output widths to agree. Suppose you\n",
    "   want a block whose output is wider than its input. Give two standard fixes\n",
    "   and the cost of each. (that section uses one of them in ResNet.)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-mxnet",
   "display_name": "d2l (MXNet)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
