{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-1 fw=all -->\n",
    "\n",
    "# Saving, Loading, and Pretrained Weights\n",
    "\n",
    "A trained network is two separate things kept in two separate places. The\n",
    "*code* is the class you wrote: its layers, its `forward` pass, the config that\n",
    "sized it. The *model state* is the tensors that training filled in: weights,\n",
    "biases, and running statistics. Optimizer moments, random-number streams, and\n",
    "the step belong to the wider *training state*. A weight file saves model state;\n",
    "a resumable checkpoint saves full training state. The code\n",
    "stays in your source repository, under version control, exactly like any other\n",
    "Python. To bring a model back you need both halves: run the code to rebuild an\n",
    "empty network, then pour the saved state into it.\n",
    "\n",
    "This split explains most of what follows. It is why a checkpoint cannot\n",
    "resurrect a model on its own, why the config object from\n",
    "that section belongs *inside* the checkpoint, and why the\n",
    "format that stores the state matters once you start sharing files with people who\n",
    "do not have your code."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-saving-loading-and-pretrained-weights"
   },
   "source": [
    "import json\n",
    "import os\n",
    "import struct\n",
    "from collections import Counter\n",
    "from dataclasses import asdict, dataclass\n",
    "from mxnet import autograd, gluon, np, npx\n",
    "from mxnet.gluon import nn\n",
    "from safetensors.numpy import load_file, save_file\n",
    "npx.set_np()"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-2 fw=all -->\n",
    "\n",
    "## State, Not Code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-3 fw=mxnet -->\n",
    "\n",
    "The state of a network is a dictionary from parameter names to arrays, the\n",
    "`collect_params` dictionary of that section. Before we save a\n",
    "whole model, the warm-up is that `npx.save` and `npx.savez` work on any array\n",
    "and on named collections of them; `npx.load` hands a saved collection back as\n",
    "a dict."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-state-not-code-1"
   },
   "source": [
    "x = np.arange(4)\n",
    "npx.savez('tensors-mx', x=x, y=np.zeros(4))\n",
    "npx.load('tensors-mx')"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'x': array([0., 1., 2., 3.]), 'y': array([0., 0., 0., 0.])}"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-4 fw=mxnet -->\n",
    "\n",
    "A model's state is one such dictionary, built for you: `collect_params` maps\n",
    "the dotted paths through the block tree (`hidden.weight`, `output.bias`) to\n",
    "the parameters, the running statistics of any BatchNorm layers included. Here\n",
    "is the tree for a small MLP."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-state-not-code-2"
   },
   "source": [
    "class MLP(nn.Block):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.hidden = nn.Dense(256, activation='relu')\n",
    "        self.output = nn.Dense(10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.output(self.hidden(x))\n",
    "\n",
    "net = MLP()\n",
    "net.initialize()\n",
    "X = np.random.uniform(size=(2, 20))\n",
    "Y = net(X)\n",
    "{name: p.shape for name, p in net.collect_params().items()}"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'hidden.weight': (256, 20),\n 'hidden.bias': (256,),\n 'output.weight': (10, 256),\n 'output.bias': (10,)}"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-5 fw=all -->\n",
    "\n",
    "Nothing in this dictionary knows it came from a class called `MLP`. The names\n",
    "and shapes refill a network with a compatible architecture and naming\n",
    "structure. Renaming an attribute or changing the nesting can break that match\n",
    "even when the computation is equivalent; the final section shows how to\n",
    "migrate such keys explicitly.\n",
    "\n",
    "## safetensors: the Interchange Format"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-6 fw=mxnet -->\n",
    "\n",
    "`npx.savez` writes MXNet's own array format: the saver accepts nothing but\n",
    "MXNet arrays, and the file holds only their names, dtypes, shapes, and raw\n",
    "bytes. There is no pickle stream and nothing to execute on load, so the\n",
    "warm-up file was already safe. The catch is reach rather than safety: hardly\n",
    "anything outside MXNet reads such a file, which matters once you hand weights\n",
    "to someone who does not run an archived framework."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-7 fw=all -->\n",
    "\n",
    "safetensors removes the problem at the root by\n",
    "having no program to run. As the figure lays out byte\n",
    "by byte, a safetensors file is an 8-byte little-endian integer giving the\n",
    "header length, then that many bytes of JSON naming each tensor's dtype, shape,\n",
    "and byte range, then the raw tensor bytes back to back. There is no opcode\n",
    "stream to interpret, so the format itself carries no executable program.\n",
    "Loaders must still validate malformed headers, offsets, and allocation sizes.\n",
    "It is also framework-neutral and memory-mappable, which is why model hubs\n",
    "commonly prefer it.\n",
    "Save and reload the MLP's state through it and confirm the round trip is exact.\n",
    "\n",
    "![The safetensors file as one horizontal byte strip: an 8-byte header length, a JSON header naming each tensor's dtype, shape, and byte offsets, and the raw tensor bytes packed back to back with no gaps, with two of the file's own data_offsets entries traced down to their exact span in the bar.](../img/bg-safetensors-layout.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-8 fw=mxnet -->\n",
    "\n",
    "There is no `safetensors.mxnet` module; the archived framework never grew a\n",
    "native binding. That turns out to be the instructive case: because the format\n",
    "is just names, dtypes, and bytes, the `safetensors.numpy` binding plus a pair\n",
    "of dict comprehensions closes the gap, `.asnumpy()` on the way out and\n",
    "`np.array` on the way back. No framework-specific machinery is involved,\n",
    "which is the format's framework-neutrality made visible."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-safetensors-the-interchange-format-3"
   },
   "source": [
    "state = {name: p.data().asnumpy() for name, p in net.collect_params().items()}\n",
    "save_file(state, 'mlp-mx.safetensors')\n",
    "restored = load_file('mlp-mx.safetensors')\n",
    "clone = MLP()\n",
    "clone.load_dict({name: np.array(v) for name, v in restored.items()})\n",
    "(clone(X) == Y).all()"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "array(True)"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-9 fw=all -->\n",
    "\n",
    "Because the header is plain JSON at a known offset, you can read it without the\n",
    "library and see there is no magic to the format."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-safetensors-the-interchange-format-4"
   },
   "source": [
    "with open('mlp-mx.safetensors', 'rb') as f:\n",
    "    n = struct.unpack('<Q', f.read(8))[0]   # header length, little-endian\n",
    "    header = json.loads(f.read(n))\n",
    "header['hidden.weight']"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'dtype': 'F32', 'shape': [256, 20], 'data_offsets': [1024, 21504]}"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-10 fw=mxnet -->\n",
    "\n",
    "`save_parameters` keeps its place for your own scratch files and checkpoints.\n",
    "safetensors is what you use to hand a model to anyone else, and for an\n",
    "archived framework it is also the exit route: current safetensors bindings\n",
    "outside MXNet can read the file you just wrote."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-11 fw=all -->\n",
    "\n",
    "## Checkpointing a Training Run\n",
    "\n",
    "A checkpoint you can resume from holds more than weights. Resuming means picking\n",
    "up the optimizer where it stopped, and Adam's state is the running first and\n",
    "second moments of the gradients from that section. Drop them and\n",
    "the optimizer restarts its momentum from zero, so the first steps after a resume\n",
    "no longer behave like a continuation. A full checkpoint therefore bundles the\n",
    "model state, the optimizer state, the RNG state (so data shuffling and dropout\n",
    "continue the same stream), the step counter, and the config that sizes the model\n",
    "when you rebuild it. the figure pairs each of those\n",
    "five compartments with the exact thing it restores on resume.\n",
    "\n",
    "![A checkpoint file's five compartments, each paired by an arrow with what it restores on resume: model state_dict with weights, optimizer state with momentum and second moments, RNG state with data order and dropout, step with schedule position, and config with architecture.](../img/bg-checkpoint-contents.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-12 fw=mxnet -->\n",
    "\n",
    "In MXNet the bundle is three files rather than one. `save_parameters` covers\n",
    "the model; the optimizer state lives with the `gluon.Trainer`, whose\n",
    "`save_states` writes Adam's moments (though not per-parameter attributes such\n",
    "as `lr_mult`, which come from your code on rebuild), and writes them with\n",
    "pickle, so a trainer-states file deserves the same your-own-disk-only caution\n",
    "as any pickle. The step counter and config travel in a JSON sidecar. One\n",
    "compartment of the figure stays empty: MXNet has no\n",
    "API to snapshot its random-number generators, only `npx.random.seed` to\n",
    "restart them, so a resumed run reseeds rather than continuing the old stream.\n",
    "The helper atomically replaces each component, but three renames are not one\n",
    "transaction. Production code should write a numbered prefix and update one\n",
    "`LATEST` manifest only after all three files are complete."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-checkpointing-a-training-run-1"
   },
   "source": [
    "def save_checkpoint(prefix, model, trainer, step, cfg=None):\n",
    "    \"\"\"Replace each component of a resumable MXNet checkpoint.\"\"\"\n",
    "    meta = {'step': step}\n",
    "    if cfg is not None:\n",
    "        meta['cfg'] = asdict(cfg)\n",
    "    model.save_parameters(prefix + '.params.tmp')\n",
    "    trainer.save_states(prefix + '.states.tmp')\n",
    "    with open(prefix + '.json.tmp', 'w') as f:\n",
    "        json.dump(meta, f)\n",
    "    for ext in ('.params', '.states', '.json'):\n",
    "        os.replace(prefix + ext + '.tmp', prefix + ext)\n",
    "\n",
    "def load_checkpoint(prefix, model, trainer=None):\n",
    "    \"\"\"Restore the state written by save_checkpoint; return the metadata.\"\"\"\n",
    "    model.load_parameters(prefix + '.params')\n",
    "    if trainer is not None:\n",
    "        trainer.load_states(prefix + '.states')\n",
    "    with open(prefix + '.json') as f:\n",
    "        return json.load(f)"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-13 fw=all -->\n",
    "\n",
    "Train a tiny regressor for a hundred steps and checkpoint it. The `Config`\n",
    "dataclass is what a rebuild reads to size the model, so it travels with the\n",
    "weights."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-checkpointing-a-training-run-2"
   },
   "source": [
    "@dataclass\n",
    "class Config:\n",
    "    in_dim: int = 20\n",
    "    hidden: int = 64\n",
    "    lr: float = 0.05\n",
    "\n",
    "def build(cfg):\n",
    "    net = nn.Sequential()\n",
    "    net.add(nn.Dense(cfg.hidden, activation='relu'), nn.Dense(1))\n",
    "    net.initialize()\n",
    "    return net\n",
    "\n",
    "npx.random.seed(1)\n",
    "cfg = Config()\n",
    "data = np.random.normal(size=(256, cfg.in_dim))\n",
    "target = data @ np.random.normal(size=(cfg.in_dim, 1)) \\\n",
    "    + 0.1 * np.random.normal(size=(256, 1))\n",
    "\n",
    "def loss(model):\n",
    "    return float(((model(data) - target) ** 2).mean())\n",
    "\n",
    "def step(model, trainer):\n",
    "    with autograd.record():\n",
    "        l = ((model(data) - target) ** 2).mean()\n",
    "    l.backward()\n",
    "    trainer.step(batch_size=1)      # the loss is already a mean\n",
    "    return float(l)\n",
    "\n",
    "net = build(cfg)\n",
    "trainer = gluon.Trainer(net.collect_params(), 'adam',\n",
    "                        {'learning_rate': cfg.lr})\n",
    "for _ in range(100):\n",
    "    step(net, trainer)\n",
    "save_checkpoint('run-mx', net, trainer, step=100, cfg=cfg)\n",
    "round(loss(net), 4)"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "0.0018"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-14 fw=all -->\n",
    "\n",
    "The restore is exact. Corrupt every parameter, load the checkpoint back, and the\n",
    "loss returns to where it was."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-15 fw=mxnet -->\n",
    "\n",
    "Note the shape of the restore: `load_parameters` and `load_states` patch the\n",
    "model and trainer in place, matching parameters by the same dotted names the\n",
    "file stores. There is no template to fill and no status object to check;\n",
    "what checking you want, you write yourself, as the final section of this\n",
    "page does with its key-set diff."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-checkpointing-a-training-run-3"
   },
   "source": [
    "for p in net.collect_params().values():\n",
    "    p.set_data(p.data() + 1.0)            # wreck the weights\n",
    "before = loss(net)\n",
    "load_checkpoint('run-mx', net, trainer)\n",
    "after = loss(net)\n",
    "f'perturbed {before:.2f} -> restored {after:.4f}'"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'perturbed 49008.15 -> restored 0.0018'"
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-16 fw=all -->\n",
    "\n",
    "Now the reason the optimizer state is in the file. Resume the run two ways from\n",
    "the same checkpoint: once restoring the optimizer, once with a fresh one holding\n",
    "only the weights. The network is near its minimum, so the correct continuation\n",
    "barely moves. A fresh Adam, with its moment estimates reset and its bias\n",
    "correction starting over, takes an oversized first step and overshoots."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-checkpointing-a-training-run-4"
   },
   "source": [
    "net_full = build(cfg)\n",
    "trainer_full = gluon.Trainer(net_full.collect_params(), 'adam',\n",
    "                             {'learning_rate': cfg.lr})\n",
    "load_checkpoint('run-mx', net_full, trainer_full)   # weights + optimizer\n",
    "full = [round(step(net_full, trainer_full), 4) for _ in range(5)]\n",
    "\n",
    "net_fresh = build(cfg)\n",
    "load_checkpoint('run-mx', net_fresh, trainer=None)  # weights only\n",
    "trainer_fresh = gluon.Trainer(net_fresh.collect_params(), 'adam',\n",
    "                              {'learning_rate': cfg.lr})\n",
    "fresh = [round(step(net_fresh, trainer_fresh), 4) for _ in range(5)]\n",
    "\n",
    "print('full  optimizer:', full)\n",
    "print('fresh optimizer:', fresh)"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "full  optimizer: [0.0018, 0.0018, 0.0017, 0.0017, 0.0017]\nfresh optimizer: [0.0018, 3.6124, 0.1789, 1.1589, 1.8473]\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-17 fw=all -->\n",
    "\n",
    "The full-state run keeps descending; the weights-only run spikes and has to claw\n",
    "its way back. That transient is the cost of forgetting the optimizer, and it is\n",
    "why \"just the weights\" is not a resumable checkpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-18 fw=mxnet -->\n",
    "\n",
    "For models too large to hold in memory, MXNet offers nothing comparable: a\n",
    "`.params` file loads whole, and the project was archived before sharded\n",
    "checkpoints and memory-mapped loading became routine. Models at that scale\n",
    "are where you leave the framework; that section returns to the\n",
    "machinery the others provide."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-19 fw=all -->\n",
    "\n",
    "## Loading Weights You Did Not Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-20 fw=mxnet -->\n",
    "\n",
    "The most common reason to load parameters is that someone else produced them.\n",
    "You take a network trained on a large dataset and adapt it: keep the learned\n",
    "feature extractor, replace the final layer for your own labels.\n",
    "`gluon.model_zoo.vision` is the built-in zoo; `pretrained=True` downloads the\n",
    "matching parameters the first time. One caveat before you rely on it: the\n",
    "download comes from the archived project's file hosting, which still worked\n",
    "when this notebook last ran but carries no promise of staying up, so\n",
    "re-verify it before building anything on top."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-loading-weights-you-did-not-train-1"
   },
   "source": [
    "net = gluon.model_zoo.vision.resnet18_v2(pretrained=True)  # ~42 MB on first run\n",
    "net.output = nn.Dense(10, in_units=512)     # new 10-class head\n",
    "net.output.initialize()\n",
    "net.output"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Dense(512 -> 10, linear)"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-21 fw=mxnet -->\n",
    "\n",
    "The partial-loading flags exist, but they report nothing: `allow_missing=True`\n",
    "skips parameters the file did not fill and `ignore_extra=True` skips file\n",
    "entries with no home in the model, both silently. So the report is yours to\n",
    "compute, the same hand-rolled key-set diff the jax tab writes. Stage a\n",
    "damaged file first: take the pretrained parameters, drop the head we just\n",
    "replaced and, to show what damage looks like, the deepest residual stage\n",
    "(`features.8`), and save what remains as a plain parameter dict."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-loading-weights-you-did-not-train-2"
   },
   "source": [
    "donor = gluon.model_zoo.vision.resnet18_v2(pretrained=True)\n",
    "file_params = {name: p.data() for name, p in donor.collect_params().items()\n",
    "               if not name.startswith(('output', 'features.8'))}\n",
    "npx.savez('resnet18-partial.params', **file_params)\n",
    "\n",
    "file_keys = set(npx.load('resnet18-partial.params'))\n",
    "damaged_target = gluon.model_zoo.vision.resnet18_v2(classes=10)\n",
    "damaged_target.initialize()\n",
    "model_keys = set(damaged_target.collect_params())\n",
    "damaged_target.load_parameters(\n",
    "    'resnet18-partial.params', allow_missing=True,\n",
    "    ignore_extra=True)  # both skip silently\n",
    "print('missing by block:', dict(Counter('.'.join(k.split('.')[:2])\n",
    "                                        for k in sorted(model_keys - file_keys))))\n",
    "print('unexpected:', sorted(file_keys - model_keys))"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "missing by block: {'features.8': 21, 'output.bias': 1, 'output.weight': 1}\nunexpected: []\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-22 fw=mxnet -->\n",
    "\n",
    "Read the diff; nothing else will print it for you. We used a fresh target so\n",
    "unexplained missing keys retain random initialization. The two `output` entries\n",
    "under *missing* are expected: that head is new and meant to start random. The\n",
    "21 `features.8` entries are a red flag, a whole residual stage the file\n",
    "failed to deliver; had this file been your only source, those layers would\n",
    "keep whatever values they started with and the features coming out of them\n",
    "would be nonsense. *Unexpected*, empty here, would list file entries with no\n",
    "home in the model, the usual sign of a renamed layer. The rule is strict\n",
    "because the flags stay silent: name\n",
    "which keys you expect in each set and treat anything else as a bug."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-loading-weights-you-did-not-train-3"
   },
   "source": [
    "for p in net.collect_params().values():\n",
    "    p.grad_req = 'null'\n",
    "for p in net.output.collect_params().values():\n",
    "    p.grad_req = 'write'\n",
    "\n",
    "trainable = sum(p.data().size for p in net.collect_params().values()\n",
    "                if p.grad_req == 'write')\n",
    "total = sum(p.data().size for p in net.collect_params().values())\n",
    "f'{trainable} trainable of {total}'"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'5130 trainable of 11187926'"
     },
     "metadata": {},
     "execution_count": 12
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-23 fw=mxnet -->\n",
    "\n",
    "`gluon.model_zoo` is one source, frozen where the project stopped; the\n",
    "Hugging Face Hub is the ecosystem-scale one, and it distributes weights as\n",
    "safetensors, which for MXNet means the numpy bridge of the previous section\n",
    "is also your import path. This section covers *how* to load and adapt\n",
    "pretrained weights; that section covers when it helps and how\n",
    "far to unfreeze."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-24 fw=all -->\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-25 fw=mxnet -->\n",
    "\n",
    "A saved model is state, not code: a dictionary of dotted-path names to arrays\n",
    "that means something only once the code that built the network runs again.\n",
    "For your own files `save_parameters` is fine, and its format is pure array\n",
    "data with no pickle to execute; for files you share, safetensors is broadly\n",
    "readable, and the `safetensors.numpy` bridge is all MXNet needs to speak it.\n",
    "A resumable checkpoint is three versioned files whose completion should be\n",
    "published by one atomic manifest update:\n",
    "parameters, trainer states (Adam's moments, stored with pickle, so keep such\n",
    "files your own), and a JSON sidecar with the step and config; the RNG stream\n",
    "cannot be snapshotted, only reseeded. Loading someone else's weights is dict\n",
    "surgery plus `allow_missing`/`ignore_extra`, and because those flags skip\n",
    "silently, the missing/unexpected key-set diff is yours to compute and read."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-26 fw=all -->\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Even if you never deploy to another machine, name two reasons to checkpoint.\n",
    "   Then consider the atomic write: if the checkpoint were written straight to its\n",
    "   final path (delete the `os.replace` from `save_checkpoint`, the\n",
    "   rename-into-place that orbax performs, or the fresh numbered files that\n",
    "   `tf.train.Checkpoint.save` writes), describe the failure a crash mid-write\n",
    "   now causes, and why the atomic version avoids it.\n",
    "1. Read the first 8 bytes of the safetensors file you wrote for the MLP as a\n",
    "   little-endian integer, as the header cell does. How large is the JSON header\n",
    "   for the MLP, and how does it grow if you double the hidden width?\n",
    "1. Save the MLP's parameters cast to `bfloat16` and load them back into a\n",
    "   `float32` model (that section). What is lost? Is that acceptable\n",
    "   for inference? For resuming training?\n",
    "1. Take two checkpoints of the regressor 50 steps apart, average their weight\n",
    "   tensors into a third set of parameters, and evaluate it. The result previews\n",
    "   weight averaging."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-mxnet",
   "display_name": "d2l (MXNet)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
