{
 "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",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torchvision.models import resnet18, ResNet18_Weights\n",
    "from safetensors.torch import load_file, save_file\n",
    "from d2l import torch as d2l"
   ],
   "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=pytorch -->\n",
    "\n",
    "The state of a network is a dictionary from parameter names to tensors, the\n",
    "`state_dict` of that section. Before we save a whole model, the\n",
    "warm-up is that the same `save`/`load` calls work on any tensors, and on the\n",
    "lists and dicts that hold them."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-state-not-code-1"
   },
   "source": [
    "x = torch.arange(4)\n",
    "torch.save({'x': x, 'y': torch.zeros(4)}, 'tensors.pt')\n",
    "torch.load('tensors.pt', weights_only=True)"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'x': tensor([0, 1, 2, 3]), 'y': tensor([0., 0., 0., 0.])}"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-4 fw=pytorch -->\n",
    "\n",
    "A model's `state_dict` is one such dictionary, built for you. The keys are the\n",
    "dotted paths through the module tree (`hidden.weight`, `output.bias`); the values\n",
    "are the tensors, buffers included. Here is the tree for a small MLP."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-state-not-code-2"
   },
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.hidden = nn.LazyLinear(256)\n",
    "        self.output = nn.LazyLinear(10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.output(F.relu(self.hidden(x)))\n",
    "\n",
    "net = MLP()\n",
    "X = torch.randn(2, 20)\n",
    "Y = net(X)\n",
    "{name: tuple(t.shape) for name, t in net.state_dict().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=pytorch -->\n",
    "\n",
    "`torch.save` writes its files with Python's `pickle`, which does not store data\n",
    "so much as a program that *reconstructs* data. Unpickling runs that program. For\n",
    "a file you wrote and never let out of your control this is harmless. For a file\n",
    "you downloaded it is a remote-code-execution vector: loading the weights can run\n",
    "whatever the author's pickle stream tells it to."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-safetensors-the-interchange-format-1"
   },
   "source": [
    "class Tripwire:\n",
    "    def __reduce__(self):\n",
    "        return (print, ('*** payload executed while loading ***',))\n",
    "\n",
    "torch.save(Tripwire(), 'tripwire.pt')\n",
    "_ = torch.load('tripwire.pt', weights_only=False)  # the payload runs here"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "*** payload executed while loading ***\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-7 fw=pytorch -->\n",
    "\n",
    "Since version 2.6, PyTorch defaults `torch.load` to `weights_only=True`, which\n",
    "refuses any pickle opcode that is not a plain tensor. The same file is now\n",
    "rejected instead of executed."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-safetensors-the-interchange-format-2"
   },
   "source": [
    "try:\n",
    "    torch.load('tripwire.pt', weights_only=True)  # the default in this torch\n",
    "except Exception as e:\n",
    "    print(type(e).__name__, str(e).splitlines()[0])"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "UnpicklingError Weights only load failed. This file can still be loaded, to do so you have two options, \u001b[1mdo those steps only if you trust the source of the checkpoint\u001b[0m. \n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-8 fw=pytorch -->\n",
    "\n",
    "The allowlist behind `weights_only=True` is defense in depth, not a sandbox: it\n",
    "has itself had bypasses patched."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-9 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": "code",
   "metadata": {
    "id": "saving-loading-safetensors-the-interchange-format-3"
   },
   "source": [
    "save_file(net.state_dict(), 'mlp.safetensors')\n",
    "clone = MLP()\n",
    "clone(X)                                   # materialize the lazy layers first\n",
    "clone.load_state_dict(load_file('mlp.safetensors'))\n",
    "clone.eval()\n",
    "torch.equal(clone(X), Y)"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-10 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.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": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'dtype': 'F32', 'shape': [256, 20], 'data_offsets': [1024, 21504]}"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-11 fw=pytorch -->\n",
    "\n",
    "`torch.save` keeps its place for your own scratch files and for the older code\n",
    "you will still meet. safetensors is what you use to hand a model to anyone else."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-12 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-13 fw=pytorch -->\n",
    "\n",
    "Two details separate a checkpoint from a corrupted file. First, keep the contents\n",
    "to tensors and primitives so the file loads under `weights_only=True`; a\n",
    "dataclass config goes in as a plain dict via `asdict`. Second, write atomically:\n",
    "save to a temporary path and `os.replace` it into place, so a crash mid-write\n",
    "leaves the previous good checkpoint untouched rather than a half-written one."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-checkpointing-a-training-run-1"
   },
   "source": [
    "def save_checkpoint(path, model, optimizer, step, cfg=None):\n",
    "    \"\"\"Atomically write a resumable training checkpoint.\"\"\"\n",
    "    ckpt = {'model': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'step': step,\n",
    "            'cpu_rng': torch.get_rng_state()}\n",
    "    if cfg is not None:\n",
    "        ckpt['cfg'] = asdict(cfg)\n",
    "    if torch.cuda.is_available():\n",
    "        ckpt['cuda_rng'] = torch.cuda.get_rng_state_all()\n",
    "    tmp = path + '.tmp'\n",
    "    torch.save(ckpt, tmp)\n",
    "    os.replace(tmp, path)\n",
    "\n",
    "def load_checkpoint(path, model, optimizer=None):\n",
    "    \"\"\"Restore the state written by save_checkpoint; return the raw dict.\"\"\"\n",
    "    ckpt = torch.load(path, weights_only=True)\n",
    "    model.load_state_dict(ckpt['model'])\n",
    "    if optimizer is not None:\n",
    "        optimizer.load_state_dict(ckpt['optimizer'])\n",
    "    torch.set_rng_state(ckpt['cpu_rng'])\n",
    "    if torch.cuda.is_available() and 'cuda_rng' in ckpt:\n",
    "        torch.cuda.set_rng_state_all(ckpt['cuda_rng'])\n",
    "    return ckpt"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-14 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",
    "    return nn.Sequential(nn.Linear(cfg.in_dim, cfg.hidden), nn.ReLU(),\n",
    "                         nn.Linear(cfg.hidden, 1))\n",
    "\n",
    "torch.manual_seed(1)\n",
    "cfg = Config()\n",
    "data = torch.randn(256, cfg.in_dim)\n",
    "target = data @ torch.randn(cfg.in_dim, 1) + 0.1 * torch.randn(256, 1)\n",
    "loss = nn.MSELoss()\n",
    "\n",
    "def step(model, opt):\n",
    "    opt.zero_grad()\n",
    "    l = loss(model(data), target)\n",
    "    l.backward()\n",
    "    opt.step()\n",
    "    return l.item()\n",
    "\n",
    "net = build(cfg)\n",
    "opt = torch.optim.Adam(net.parameters(), lr=cfg.lr)\n",
    "for _ in range(100):\n",
    "    step(net, opt)\n",
    "save_checkpoint('run.pt', net, opt, step=100, cfg=cfg)\n",
    "round(loss(net(data), target).item(), 4)"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "0.0016"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-15 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": "code",
   "metadata": {
    "id": "saving-loading-checkpointing-a-training-run-3"
   },
   "source": [
    "with torch.no_grad():\n",
    "    for p in net.parameters():\n",
    "        p.add_(1.0)                       # wreck the weights\n",
    "before = loss(net(data), target).item()\n",
    "load_checkpoint('run.pt', net, opt)\n",
    "after = loss(net(data), target).item()\n",
    "f'perturbed {before:.2f} -> restored {after:.4f}'"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'perturbed 47291.83 -> restored 0.0016'"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "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",
    "opt_full = torch.optim.Adam(net_full.parameters(), lr=cfg.lr)\n",
    "load_checkpoint('run.pt', net_full, opt_full)          # weights + optimizer\n",
    "full = [round(step(net_full, opt_full), 4) for _ in range(5)]\n",
    "\n",
    "net_fresh = build(cfg)\n",
    "load_checkpoint('run.pt', net_fresh, optimizer=None)   # weights only\n",
    "opt_fresh = torch.optim.Adam(net_fresh.parameters(), lr=cfg.lr)\n",
    "fresh = [round(step(net_fresh, opt_fresh), 4) for _ in range(5)]\n",
    "\n",
    "print('full  optimizer:', full)\n",
    "print('fresh optimizer:', fresh)"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "full  optimizer: [0.0016, 0.0016, 0.0015, 0.0015, 0.0015]\nfresh optimizer: [0.0016, 3.8324, 0.2922, 1.0421, 1.845]\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=pytorch -->\n",
    "\n",
    "For models too large to hold in memory, checkpoints are split across several\n",
    "files with an index, and `torch.load(..., mmap=True)` pages tensors off disk on\n",
    "demand instead of copying the whole file up front. Combined with meta-device\n",
    "construction and `load_state_dict(..., assign=True)`, this loads such a model\n",
    "without ever allocating its randomly-initialized weights;\n",
    "that section returns to the machinery when models get that big."
   ]
  },
  {
   "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=pytorch -->\n",
    "\n",
    "The most common reason to load a `state_dict` is that someone else produced it.\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. The mechanics are\n",
    "`state_dict` manipulation. torchvision serves the weights through a `weights=`\n",
    "enum, which also downloads the matching parameters the first time."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-loading-weights-you-did-not-train-1"
   },
   "source": [
    "net = resnet18(weights=ResNet18_Weights.DEFAULT)   # ~45 MB on first run\n",
    "net.fc = nn.Linear(net.fc.in_features, 10)          # new 10-class head\n",
    "net.fc"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Linear(in_features=512, out_features=10, bias=True)"
     },
     "metadata": {},
     "execution_count": 12
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-21 fw=pytorch -->\n",
    "\n",
    "A `state_dict` is an ordinary Python dict, so adapting one is ordinary dict\n",
    "surgery. We drop the pretrained 1000-class head (we just replaced it) and, to\n",
    "show what a damaged file looks like, also drop one residual block. Loading with\n",
    "`strict=False` then returns a report of what did not line up instead of raising."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-loading-weights-you-did-not-train-2"
   },
   "source": [
    "pretrained = ResNet18_Weights.DEFAULT.get_state_dict(progress=False)\n",
    "pretrained = {k: v for k, v in pretrained.items()\n",
    "              if not k.startswith('fc.') and not k.startswith('layer4.')}\n",
    "damaged_target = resnet18(weights=None)\n",
    "damaged_target.fc = nn.Linear(damaged_target.fc.in_features, 10)\n",
    "report = damaged_target.load_state_dict(pretrained, strict=False)\n",
    "print('missing by block:', dict(Counter(k.split('.')[0]\n",
    "                                         for k in report.missing_keys)))\n",
    "print('unexpected:', report.unexpected_keys)"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "missing by block: {'layer4': 25, 'fc': 2}\nunexpected: []\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-22 fw=pytorch -->\n",
    "\n",
    "Read this report; do not discard it. We loaded the damaged file into a fresh\n",
    "target so every unexplained missing key really remains random. `missing_keys` lists parameters the model\n",
    "has but the file did not fill. The two `fc` entries are expected: that head is\n",
    "new and meant to start random. The `layer4` entries are a red flag, a whole block\n",
    "of the backbone left uninitialized, which here means the incoming file was\n",
    "incomplete and would produce nonsense features. `unexpected_keys`, empty here,\n",
    "would list names in the file with no home in the model, the usual sign of a\n",
    "renamed layer. The rule is to name which keys you expect to be missing and treat\n",
    "anything else as a bug."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-loading-weights-you-did-not-train-3"
   },
   "source": [
    "for p in net.parameters():\n",
    "    p.requires_grad = False\n",
    "for p in net.fc.parameters():\n",
    "    p.requires_grad = True\n",
    "\n",
    "trainable = sum(p.numel() for p in net.parameters() if p.requires_grad)\n",
    "total = sum(p.numel() for p in net.parameters())\n",
    "f'{trainable} trainable of {total}'"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'5130 trainable of 11181642'"
     },
     "metadata": {},
     "execution_count": 14
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-23 fw=pytorch -->\n",
    "\n",
    "torchvision is one source; the Hugging Face Hub is the ecosystem-scale one, and\n",
    "it distributes its weights as safetensors, which closes the loop with the format\n",
    "of the previous section. This section covers *how* to load and adapt pretrained\n",
    "weights; that section covers when it helps and how 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=pytorch -->\n",
    "\n",
    "A saved model is state, not code: a `state_dict` of tensors that means something\n",
    "only once the code that built the network runs again. For your own files\n",
    "`torch.save` is fine; for files you share, safetensors stores the same tensors\n",
    "with no executable pickle, which is why hubs standardize on it. A resumable\n",
    "checkpoint bundles more than weights: optimizer state, RNG state, step, and\n",
    "config, written atomically, or a resume restarts the optimizer's momentum from\n",
    "zero. Loading someone else's weights is dict surgery plus `strict=False`, and the\n",
    "missing/unexpected report is a diagnostic to read rather than a warning to\n",
    "silence."
   ]
  },
  {
   "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-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
