{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8b4cadb2",
   "metadata": {
    "origin_pos": 1
   },
   "source": [
    "# File I/O\n",
    "\n",
    "So far we have discussed how to process data and how\n",
    "to build, train, and test deep learning models.\n",
    "However, at some point we will hopefully be happy enough\n",
    "with the learned models that we will want\n",
    "to save the results for later use in various contexts\n",
    "(perhaps even to make predictions in deployment).\n",
    "Additionally, when running a long training process,\n",
    "the best practice is to periodically save intermediate results (checkpointing)\n",
    "to ensure that we do not lose several days' worth of computation\n",
    "if we trip over the power cord of our server.\n",
    "Thus it is time to learn how to load and store\n",
    "both individual weight vectors and entire models.\n",
    "This section addresses both issues.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3dff9f92",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T20:14:02.401475Z",
     "iopub.status.busy": "2023-08-18T20:14:02.401201Z",
     "iopub.status.idle": "2023-08-18T20:14:04.212167Z",
     "shell.execute_reply": "2023-08-18T20:14:04.211057Z"
    },
    "origin_pos": 3,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71d2d6c9",
   "metadata": {
    "origin_pos": 6
   },
   "source": [
    "## (**Loading and Saving Tensors**)\n",
    "\n",
    "For individual tensors, we can directly\n",
    "invoke the `load` and `save` functions\n",
    "to read and write them respectively.\n",
    "Both functions require that we supply a name,\n",
    "and `save` requires as input the variable to be saved.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "41c45edb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T20:14:04.216976Z",
     "iopub.status.busy": "2023-08-18T20:14:04.216107Z",
     "iopub.status.idle": "2023-08-18T20:14:04.247988Z",
     "shell.execute_reply": "2023-08-18T20:14:04.246686Z"
    },
    "origin_pos": 8,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "x = torch.arange(4)\n",
    "torch.save(x, 'x-file')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4d4cdfe",
   "metadata": {
    "origin_pos": 11
   },
   "source": [
    "We can now read the data from the stored file back into memory.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5cb11ed0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T20:14:04.253634Z",
     "iopub.status.busy": "2023-08-18T20:14:04.253018Z",
     "iopub.status.idle": "2023-08-18T20:14:04.262485Z",
     "shell.execute_reply": "2023-08-18T20:14:04.261617Z"
    },
    "origin_pos": 13,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 1, 2, 3])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x2 = torch.load('x-file')\n",
    "x2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "431f40d8",
   "metadata": {
    "origin_pos": 16
   },
   "source": [
    "We can [**store a list of tensors and read them back into memory.**]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "86dba6ad",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T20:14:04.267689Z",
     "iopub.status.busy": "2023-08-18T20:14:04.267013Z",
     "iopub.status.idle": "2023-08-18T20:14:04.275474Z",
     "shell.execute_reply": "2023-08-18T20:14:04.274471Z"
    },
    "origin_pos": 18,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([0, 1, 2, 3]), tensor([0., 0., 0., 0.]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = torch.zeros(4)\n",
    "torch.save([x, y],'x-files')\n",
    "x2, y2 = torch.load('x-files')\n",
    "(x2, y2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e03e08f",
   "metadata": {
    "origin_pos": 21
   },
   "source": [
    "We can even [**write and read a dictionary that maps\n",
    "from strings to tensors.**]\n",
    "This is convenient when we want\n",
    "to read or write all the weights in a model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1b8f14c8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T20:14:04.279967Z",
     "iopub.status.busy": "2023-08-18T20:14:04.279484Z",
     "iopub.status.idle": "2023-08-18T20:14:04.286799Z",
     "shell.execute_reply": "2023-08-18T20:14:04.286045Z"
    },
    "origin_pos": 23,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'x': tensor([0, 1, 2, 3]), 'y': tensor([0., 0., 0., 0.])}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mydict = {'x': x, 'y': y}\n",
    "torch.save(mydict, 'mydict')\n",
    "mydict2 = torch.load('mydict')\n",
    "mydict2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1f21d5a",
   "metadata": {
    "origin_pos": 26
   },
   "source": [
    "## [**Loading and Saving Model Parameters**]\n",
    "\n",
    "Saving individual weight vectors (or other tensors) is useful,\n",
    "but it gets very tedious if we want to save\n",
    "(and later load) an entire model.\n",
    "After all, we might have hundreds of\n",
    "parameter groups sprinkled throughout.\n",
    "For this reason the deep learning framework provides built-in functionalities\n",
    "to load and save entire networks.\n",
    "An important detail to note is that this\n",
    "saves model *parameters* and not the entire model.\n",
    "For example, if we have a 3-layer MLP,\n",
    "we need to specify the architecture separately.\n",
    "The reason for this is that the models themselves can contain arbitrary code,\n",
    "hence they cannot be serialized as naturally.\n",
    "Thus, in order to reinstate a model, we need\n",
    "to generate the architecture in code\n",
    "and then load the parameters from disk.\n",
    "(**Let's start with our familiar MLP.**)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6917c6ad",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T20:14:04.290900Z",
     "iopub.status.busy": "2023-08-18T20:14:04.290420Z",
     "iopub.status.idle": "2023-08-18T20:14:04.301461Z",
     "shell.execute_reply": "2023-08-18T20:14:04.300490Z"
    },
    "origin_pos": 28,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "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(size=(2, 20))\n",
    "Y = net(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45d0c48e",
   "metadata": {
    "origin_pos": 31
   },
   "source": [
    "Next, we [**store the parameters of the model as a file**] with the name \"mlp.params\".\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "88dfe184",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T20:14:04.304908Z",
     "iopub.status.busy": "2023-08-18T20:14:04.304617Z",
     "iopub.status.idle": "2023-08-18T20:14:04.309701Z",
     "shell.execute_reply": "2023-08-18T20:14:04.308927Z"
    },
    "origin_pos": 33,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "torch.save(net.state_dict(), 'mlp.params')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ddcdcbd",
   "metadata": {
    "origin_pos": 36
   },
   "source": [
    "To recover the model, we instantiate a clone\n",
    "of the original MLP model.\n",
    "Instead of randomly initializing the model parameters,\n",
    "we [**read the parameters stored in the file directly**].\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0a8e9c03",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T20:14:04.312721Z",
     "iopub.status.busy": "2023-08-18T20:14:04.312444Z",
     "iopub.status.idle": "2023-08-18T20:14:04.320044Z",
     "shell.execute_reply": "2023-08-18T20:14:04.319202Z"
    },
    "origin_pos": 38,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MLP(\n",
       "  (hidden): LazyLinear(in_features=0, out_features=256, bias=True)\n",
       "  (output): LazyLinear(in_features=0, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clone = MLP()\n",
    "clone.load_state_dict(torch.load('mlp.params'))\n",
    "clone.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b209774",
   "metadata": {
    "origin_pos": 41
   },
   "source": [
    "Since both instances have the same model parameters,\n",
    "the computational result of the same input `X` should be the same.\n",
    "Let's verify this.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d65ae251",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T20:14:04.323094Z",
     "iopub.status.busy": "2023-08-18T20:14:04.322816Z",
     "iopub.status.idle": "2023-08-18T20:14:04.330451Z",
     "shell.execute_reply": "2023-08-18T20:14:04.329304Z"
    },
    "origin_pos": 42,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y_clone = clone(X)\n",
    "Y_clone == Y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "109a1a9f",
   "metadata": {
    "origin_pos": 44
   },
   "source": [
    "## Summary\n",
    "\n",
    "The `save` and `load` functions can be used to perform file I/O for tensor objects.\n",
    "We can save and load the entire sets of parameters for a network via a parameter dictionary.\n",
    "Saving the architecture has to be done in code rather than in parameters.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Even if there is no need to deploy trained models to a different device, what are the practical benefits of storing model parameters?\n",
    "1. Assume that we want to reuse only parts of a network to be incorporated into a network having a different architecture. How would you go about using, say the first two layers from a previous network in a new network?\n",
    "1. How would you go about saving the network architecture and parameters? What restrictions would you impose on the architecture?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "815c6198",
   "metadata": {
    "origin_pos": 46,
    "tab": [
     "pytorch"
    ]
   },
   "source": [
    "[Discussions](https://discuss.d2l.ai/t/61)\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "required_libs": []
 },
 "nbformat": 4,
 "nbformat_minor": 5
}