{
 "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 d2l import torch as d2l\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F"
   ],
   "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=pytorch -->\n",
    "\n",
    "In PyTorch the module class is `nn.Module`. We have used one of its subclasses\n",
    "all along: `nn.Sequential` builds a model from a chain of layers, here the\n",
    "familiar MLP with a 256-unit ReLU hidden layer and a 10-unit output layer."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-1"
   },
   "source": [
    "net = nn.Sequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))\n",
    "\n",
    "X = torch.rand(2, 20)\n",
    "net(X).shape"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([2, 10])"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-4 fw=pytorch -->\n",
    "\n",
    "`Sequential` is not a special construct. It is itself a module whose forward\n",
    "computation runs its children in order, and the children are the three modules\n",
    "we passed to it, stored under names in a registry:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-2"
   },
   "source": [
    "net._modules"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'0': Linear(in_features=20, out_features=256, bias=True),\n '1': ReLU(),\n '2': Linear(in_features=256, out_features=10, bias=True)}"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-5 fw=pytorch -->\n",
    "\n",
    "This registry is what the `nn.Module` machinery traverses: `net.parameters()`\n",
    "collects parameters by walking `_modules` recursively, and the same walk\n",
    "underlies device movement and serialization. A module the registry does not\n",
    "contain might as well not exist, a fact we will exploit for a demonstration\n",
    "shortly."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-3"
   },
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.hidden = nn.LazyLinear(256)\n",
    "        self.out = nn.LazyLinear(10)\n",
    "\n",
    "    def forward(self, X):\n",
    "        return self.out(F.relu(self.hidden(X)))"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-4"
   },
   "source": [
    "net = MLP()\n",
    "net(X).shape"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([2, 10])"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-6 fw=pytorch -->\n",
    "\n",
    "Two details do the work here. First, `self.hidden = nn.LazyLinear(256)` is not\n",
    "an ordinary attribute assignment: `nn.Module` intercepts `__setattr__`, sees\n",
    "that the value is a module, and adds it to the registry we just inspected. That\n",
    "is why both layers' parameters show up in `net.parameters()` 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=pytorch -->\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 under a name, and loop\n",
    "over the children in `forward`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-1"
   },
   "source": [
    "class MySequential(nn.Module):\n",
    "    def __init__(self, *args):\n",
    "        super().__init__()\n",
    "        for idx, module in enumerate(args):\n",
    "            self.add_module(str(idx), module)\n",
    "\n",
    "    def forward(self, X):\n",
    "        for module in self.children():\n",
    "            X = module(X)\n",
    "        return X"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-9 fw=pytorch -->\n",
    "\n",
    "`add_module` writes a child into the registry under a string name (that is\n",
    "where the `'0'`, `'1'`, `'2'` keys above came from), and `children()` iterates\n",
    "the registry in insertion order. Our version is a drop-in replacement:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-2"
   },
   "source": [
    "net = MySequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))\n",
    "net(X).shape"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([2, 10])"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-10 fw=pytorch -->\n",
    "\n",
    "The registration step is easy to lose. The following module looks reasonable,\n",
    "and its forward pass works, so nothing appears wrong:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-3"
   },
   "source": [
    "class PlainListMLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.layers = [nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10)]\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, sum(p.numel() for p in net.parameters())"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(torch.Size([2, 10]), 0)"
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-11 fw=pytorch -->\n",
    "\n",
    "The model computes, yet it owns zero parameters. A plain Python list is not a\n",
    "module, so the `__setattr__` interception ignores it and nothing inside it is\n",
    "registered. The layers still hold perfectly good tensors, which is why\n",
    "`forward` runs. But `net.parameters()` is empty: hand this model to an\n",
    "optimizer and training proceeds without a single error while updating nothing;\n",
    "`net.to(device)` leaves every layer behind on the CPU; `net.state_dict()`\n",
    "checkpoints an empty model. Because no step raises an exception, this bug is\n",
    "usually diagnosed by staring at a loss curve that refuses to move."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-4"
   },
   "source": [
    "class ModuleListMLP(PlainListMLP):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList(self.layers)\n",
    "\n",
    "net = ModuleListMLP()\n",
    "net(X).shape, sum(p.numel() for p in net.parameters())"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(torch.Size([2, 10]), 7946)"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-12 fw=pytorch -->\n",
    "\n",
    "Same forward pass, 7946 registered parameters. The division of labor among the\n",
    "containers is now clear. `nn.Sequential` registers its children and supplies\n",
    "the run-them-in-order `forward`. `nn.ModuleList` registers a list of children\n",
    "but supplies no `forward` at all; you keep writing the loop, which is exactly\n",
    "what you want when the loop body is not a plain chain (blocks that take extra\n",
    "arguments, or a skip connection around each block). `nn.ModuleDict` does the\n",
    "same for children indexed by name. Transformer implementations conventionally\n",
    "keep their stack of blocks in an `nn.ModuleList` and their named parts\n",
    "(embedding, final normalization, output head) as 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=pytorch -->\n",
    "\n",
    "`forward` is an ordinary Python method. Nothing restricts it to chaining\n",
    "children: it can branch, loop, call any tensor function, and combine\n",
    "intermediate results however it likes. The loop in `ModuleListMLP` already used\n",
    "this freedom. Its most consequential one-line use is the *residual\n",
    "connection*, the wiring idiom at the heart of ResNets and Transformers alike:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-forward-is-just-python-1"
   },
   "source": [
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, num_hiddens):\n",
    "        super().__init__()\n",
    "        self.body = nn.Sequential(\n",
    "            nn.Linear(num_hiddens, num_hiddens), nn.ReLU(),\n",
    "            nn.Linear(num_hiddens, 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=pytorch -->\n",
    "\n",
    "That is why we gave `body` explicit `nn.Linear` layers rather than lazy ones."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-forward-is-just-python-2"
   },
   "source": [
    "block = ResidualBlock(24)\n",
    "block(torch.randn(2, 24)).shape"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([2, 24])"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-17 fw=pytorch -->\n",
    "\n",
    "`forward` may also use state that is neither an input nor a parameter. Suppose\n",
    "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 = torch.tensor(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.state_dict(), list(block.state_dict())[: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=pytorch -->\n",
    "\n",
    "`alpha` enters the computation, but it is not a parameter: it never appears in\n",
    "`named_parameters()`, so the optimizer never touches it. That much we wanted.\n",
    "Storing it as a plain attribute has a cost we did not want, though: as the\n",
    "output shows, it is missing from `state_dict()` as well, so it will not be\n",
    "saved with the model, and `.to(device)` will not move it. Some state is not a\n",
    "parameter but must still travel with the model; the registered home for such\n",
    "state is a *buffer*, 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=pytorch -->\n",
    "\n",
    "We have been doing something odd since our first MLP without commenting on it:\n",
    "`nn.LazyLinear(256)` names only the layer's *output* width. Its weight matrix\n",
    "has shape `(256, in_features)`, and we never said what `in_features` is.\n",
    "The layer cannot know it at construction time, since it depends on the data it\n",
    "will receive. So it does not allocate parameters at construction time at all:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-lazy-initialization-shapes-from-data-1"
   },
   "source": [
    "net = nn.Sequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))\n",
    "net[0].weight"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "<UninitializedParameter>"
     },
     "metadata": {},
     "execution_count": 13
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-21 fw=pytorch -->\n",
    "\n",
    "The weight is a placeholder. The first time data flows through, the layer\n",
    "reads the input width from the batch, allocates and initializes a real weight,\n",
    "and replaces itself with a plain `nn.Linear`. Its output width then fixes the\n",
    "input of the next lazy layer, and shapes cascade through the whole model:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-lazy-initialization-shapes-from-data-2"
   },
   "source": [
    "net(X)\n",
    "net[0].weight.shape"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([256, 20])"
     },
     "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=pytorch -->\n",
    "\n",
    "Before the first forward pass, lazy layers contain registered\n",
    "`UninitializedParameter` placeholders rather than shaped arrays. An optimizer\n",
    "can be constructed over those placeholders, but reading shapes, counting\n",
    "elements, or applying a shape-dependent initializer requires a *dry run* on a\n",
    "representative batch. A related subtlety:\n",
    "the random initialization now happens at first call rather than at\n",
    "construction, so any random numbers your program draws in between shift the\n",
    "generator's state, and a fixed seed can yield different weights than the\n",
    "explicitly shaped version of the same model (that section returns to\n",
    "seeding). NNX has no lazy mode: its linear layers require both widths and\n",
    "create parameters in the constructor. PyTorch's lazy layers instead let the\n",
    "first real batch supply the missing width."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-lazy-initialization-shapes-from-data-3"
   },
   "source": [
    "@d2l.add_to_class(d2l.Module)\n",
    "def apply_init(self, inputs, init=None):\n",
    "    self.forward(*inputs)\n",
    "    if init is not None:\n",
    "        self.net.apply(init)"
   ],
   "execution_count": 15,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-24 fw=pytorch -->\n",
    "\n",
    "`nn.Module.apply(fn)` calls `fn` on every module in the tree, children first.\n",
    "It is the standard way to push a policy across an arbitrary model, one more\n",
    "operation that is a tree walk, and from that section on the idiom\n",
    "`model.apply_init([X], init)` opens most of our training scripts. A small\n",
    "demonstration:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-lazy-initialization-shapes-from-data-4"
   },
   "source": [
    "class TinyMLP(d2l.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(nn.LazyLinear(256), nn.ReLU(),\n",
    "                                 nn.LazyLinear(10))\n",
    "\n",
    "def init_xavier(module):\n",
    "    if type(module) == nn.Linear:\n",
    "        nn.init.xavier_uniform_(module.weight)\n",
    "\n",
    "model = TinyMLP()\n",
    "model.apply_init([X], init_xavier)\n",
    "model.net[0].weight.shape"
   ],
   "execution_count": 16,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([256, 20])"
     },
     "metadata": {},
     "execution_count": 16
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-25 fw=pytorch -->\n",
    "\n",
    "The dry run inside `apply_init` turned every lazy layer into a real one, after\n",
    "which `init_xavier` could match on `nn.Linear` and rewrite its weights. Which\n",
    "initializer to apply, and why Xavier's variance rule\n",
    "(that section) is a sensible default, is the subject of\n",
    "that section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-26 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.Module:\n",
    "    blocks = [ResidualBlock(cfg.d_hidden) for _ in range(cfg.num_blocks)]\n",
    "    return nn.Sequential(nn.Linear(cfg.d_in, cfg.d_hidden),\n",
    "                         *blocks, nn.Linear(cfg.d_hidden, cfg.d_out))"
   ],
   "execution_count": 17,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-27 fw=pytorch -->\n",
    "\n",
    "One config produces one architecture: an input projection into the hidden\n",
    "width, `num_blocks` identical residual blocks, and an output projection. The\n",
    "list comprehension producing `blocks` is a plain Python list, which is safe\n",
    "here for the reason that section taught:\n",
    "unpacking it into `nn.Sequential` is what registers each block. Every layer\n",
    "gets explicit widths, and this is where explicit shapes beat lazy ones: the\n",
    "config already knows every width, so explicit construction yields a fully\n",
    "materialized model with no dry run needed. Printing the model displays the\n",
    "module tree:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-2"
   },
   "source": [
    "net = build(MLPConfig())\n",
    "net"
   ],
   "execution_count": 18,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Sequential(\n  (0): Linear(in_features=784, out_features=256, bias=True)\n  (1): ResidualBlock(\n    (body): Sequential(\n      (0): Linear(in_features=256, out_features=256, bias=True)\n      (1): ReLU()\n      (2): Linear(in_features=256, out_features=256, bias=True)\n    )\n  )\n  (2): ResidualBlock(\n    (body): Sequential(\n      (0): Linear(in_features=256, out_features=256, bias=True)\n      (1): ReLU()\n      (2): Linear(in_features=256, out_features=256, bias=True)\n    )\n  )\n  (3): ResidualBlock(\n    (body): Sequential(\n      (0): Linear(in_features=256, out_features=256, bias=True)\n      (1): ReLU()\n      (2): Linear(in_features=256, out_features=256, bias=True)\n    )\n  )\n  (4): ResidualBlock(\n    (body): Sequential(\n      (0): Linear(in_features=256, out_features=256, bias=True)\n      (1): ReLU()\n      (2): Linear(in_features=256, out_features=256, bias=True)\n    )\n  )\n  (5): Linear(in_features=256, out_features=10, bias=True)\n)"
     },
     "metadata": {},
     "execution_count": 18
    }
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-3"
   },
   "source": [
    "net(torch.rand(2, 784)).shape"
   ],
   "execution_count": 19,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "torch.Size([2, 10])"
     },
     "metadata": {},
     "execution_count": 19
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-28 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",
    "    n = sum(p.numel() for p in build(cfg).parameters())\n",
    "    print(f'd_hidden={cfg.d_hidden}, num_blocks={cfg.num_blocks}: '\n",
    "          f'{n:,} parameters')"
   ],
   "execution_count": 20,
   "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-29 fw=pytorch -->\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-30 fw=all -->\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-31 fw=pytorch -->\n",
    "\n",
    "A module owns parameters, child modules, 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",
    "device movement, and serialization traverse that object graph. Children are\n",
    "discovered through registration: attribute\n",
    "assignment and the containers `nn.Sequential`, `nn.ModuleList`, and\n",
    "`nn.ModuleDict` register children, while a plain Python list hides them,\n",
    "yielding a model that runs but trains nothing. `forward` is ordinary Python;\n",
    "a residual connection is one line in it. Lazy layers declare output widths and\n",
    "infer input widths on the first forward pass, so initialization and inspection\n",
    "follow a dry run (`apply_init`). Configs turn architecture into data: a\n",
    "`dataclass` of widths and depths plus a `build` function that stacks blocks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-32 fw=all -->\n",
    "\n",
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-33 fw=pytorch -->\n",
    "\n",
    "1. Take `PlainListMLP` and catalog everything that breaks besides the empty\n",
    "   parameter list. Check `net.state_dict()`, the effect of\n",
    "   `net.to(torch.float64)` on the hidden layers' dtypes, and whether\n",
    "   `net.eval()` reaches the children. Explain how each failure follows from\n",
    "   the same missing registration.\n",
    "1. Implement a `ParallelBlock` that takes two child modules `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-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
