11_Pytorch_Advanced_Usage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "34577014-45ed-4d08-a00d-5b9f2913025f",
   "metadata": {},
   "source": [
    "- Distinguish between PyTorch classes such as nn.Module and nn.Functional, and determine their appropriate usage.\n",
    "- Customize training options, such as setting different learning rates for various layers and initializing weights.\n",
    "- Utilize Scheduling with PyTorch.\n",
    "- Saving and Loading your Model\n",
    "\n",
    "#### nn.Functional vs nn.Module\n",
    "\n",
    "The use of torch.nn.Module and torch.nn.functional in PyTorch often arises, especially when reading open-source code. Both are essential components for building neural network architectures, but they serve different purposes.\n",
    "\n",
    "`torch.nn.Module` serves as the cornerstone of PyTorch. It allows users to define complex neural network architectures by creating custom modules. These modules encapsulate parameters and operations, making it easy to organize and reuse code. When using `torch.nn.Module`, you define a module object and then invoke its forward method to execute it. This approach follows an object-oriented paradigm.\n",
    "\n",
    "On the other hand, `torch.nn.functional` provides a set of pre-defined layers and activation functions that can be directly applied to input tensors. These functions operate on tensors and are stateless, meaning they do not contain any parameters themselves. For example, to perform operations like rescaling an image tensor, you can call functions like `torch.nn.functional.interpolate` directly on the input tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e5395a7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True' # Required for matplotlib crash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6c890295-2196-438b-9006-d63ec3c5b432",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "inp = torch.randn(1,3,64,64)     # random input image\n",
    "\n",
    "# torch.nn\n",
    "avg_pool = nn.AvgPool2d(4)     # create an object\n",
    "nn_out = avg_pool(inp)         # invoke the forward method\n",
    "\n",
    "# torch.nn.Functional\n",
    "f_out = F.avg_pool2d(inp, 4)\n",
    "\n",
    "print (torch.equal(nn_out, f_out))  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "383cc97f-4377-4be6-a405-7967e9b6ae48",
   "metadata": {},
   "source": [
    "So how do we choose what to use when? When the layer / activation / loss we are implementing has a loss.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4873c281-6cfe-41c7-9f40-6e624ae47b89",
   "metadata": {},
   "source": [
    "#### Understanding the Stateful Nature of Layers\n",
    "\n",
    "A layer in a neural network can be viewed as a function that transforms input data into output data through mathematical operations such as convolutions or matrix multiplications. However, layers in deep learning models also have parameters, or weights, that need to be learned during training. These weights are updated iteratively through optimization algorithms like gradient descent.\n",
    "\n",
    "The crucial point to understand is that these weights are not fixed; they change as we train the model. This means that layers in neural networks have a state that evolves over time as we update the weights based on the training data.\n",
    "\n",
    "To implement a layer as a function, we would need to handle the storage and update of these weights separately from the function itself. This introduces complexity and can be cumbersome to manage, especially as the number of layers and parameters in the model grows.\n",
    "\n",
    "Alternatively, we can encapsulate both the function and its associated state (i.e., the weights) within a class. This approach simplifies the implementation by keeping the stateful variables within the class itself, eliminating the need to manage external data structures. In PyTorch, we achieve this by defining a class that inherits from `torch.nn.Module`, where each layer is implemented as a member function of the class.\n",
    "\n",
    "Using `torch.nn.Module` objects is preferred when a layer requires weights or other stateful variables that influence its behavior during training. Examples include layers like dropout or batch normalization, which behave differently during training and inference.\n",
    "\n",
    "On the other hand, for operations that do not require state or weights, such as resizing or average pooling, we can use the corresponding functions from `torch.nn.functional`. These functions operate directly on input tensors without maintaining any internal state.\n",
    "\n",
    "Although most `torch.nn.Module` classes have equivalent functions in `torch.nn.functional`, it's essential to consider the stateful nature of layers during practical work. Choosing between `torch.nn.Module` and `torch.nn.functional` based on the presence of state or weights helps ensure clear and efficient implementation of neural network architectures."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b7fa9cd-68cb-429b-97f9-9158aefcb97f",
   "metadata": {},
   "source": [
    "#### nn.Parameter\n",
    "\n",
    "An important class in PyTorch is the `nn.Parameter` class, which to my surprise, has gotten little coverage in PyTorch introductory texts. Consider the following case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6fe5ff6f-d8c7-4a36-839b-65a8772351f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Parameter containing:\n",
      "tensor([[-0.1972, -0.2504,  0.0037, -0.1255,  0.1843, -0.2034, -0.2909, -0.2230,\n",
      "          0.2202, -0.1253],\n",
      "        [-0.1081,  0.0295,  0.1450, -0.1925,  0.0685, -0.0880,  0.0645, -0.0136,\n",
      "          0.1366,  0.1008],\n",
      "        [-0.3049, -0.1564,  0.2202,  0.1279,  0.0413, -0.0116,  0.0106, -0.0782,\n",
      "          0.1009,  0.2774],\n",
      "        [-0.2828,  0.0182, -0.0097, -0.1376,  0.2549, -0.1233, -0.1226,  0.0767,\n",
      "         -0.2976,  0.1205],\n",
      "        [ 0.0766,  0.0442,  0.0079,  0.1717,  0.1839, -0.1046,  0.0254,  0.2459,\n",
      "          0.0909,  0.0334]], requires_grad=True), Parameter containing:\n",
      "tensor([ 0.0823,  0.0625, -0.1644, -0.2324,  0.0243], requires_grad=True)]\n"
     ]
    }
   ],
   "source": [
    "class net(nn.Module):\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.conv = nn.Linear(10,5)\n",
    "    \n",
    "  def forward(self, x):\n",
    "    return self.linear(x)\n",
    "\n",
    "\n",
    "myNet = net()\n",
    "print(list(myNet.parameters()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f60ac99-893e-497c-97bd-3dac51757419",
   "metadata": {},
   "source": [
    "Every `nn.Module` in PyTorch includes a `parameters()` function that returns its trainable parameters. These parameters are explicitly defined by the authors of PyTorch when they implement the module's functionality. For instance, in the definition of `nn.Conv2d`, the parameters representing the convolutional kernel weights and bias are explicitly defined.\n",
    "\n",
    "However, it's worth noting that when we define a neural network (`net`), we don't need to explicitly add the parameters of each individual layer (`nn.Conv2d`, `nn.Linear`, etc.) to the parameters of the `net`. Instead, this process happens implicitly due to the way PyTorch handles the registration of parameters.\n",
    "\n",
    "Internally, PyTorch utilizes the `nn.Parameter` class, which is a subclass of the Tensor class, to represent trainable parameters. When the `parameters()` function of an `nn.Module` object is invoked, it returns all of its members that are instances of `nn.Parameter`.\n",
    "\n",
    "Moreover, all the trainable weights within `nn.Module` classes are implemented as `nn.Parameter` objects. When one `nn.Module` (e.g., `nn.Conv2d`) is assigned as a member of another `nn.Module` (e.g., `net`), the parameters of the assigned module (i.e., the weights of `nn.Conv2d`) are automatically added to the parameters of the object to which it's being assigned (parameters of net). This process is referred to as registering the parameters of an `nn.Module`.\n",
    "\n",
    "It's important to note that if you try to assign a regular tensor to an `nn.Module` object, it won't be included in the parameters returned by the `parameters()` function unless you explicitly define it as an `nn.Parameter` object. This design choice allows for flexibility in scenarios where you might need to cache a non-differentiable tensor, such as caching previous outputs in recurrent neural networks (RNNs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5bb5d245-4a7f-420b-8426-cb7b0c0a48ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Parameter containing:\n",
      "tensor([[-0.2800, -0.2969, -0.2541, -0.2142, -0.1865,  0.0554, -0.1238,  0.2148,\n",
      "          0.1229, -0.0278],\n",
      "        [ 0.1394,  0.0280, -0.2778, -0.2146, -0.0930,  0.0186,  0.0074,  0.3147,\n",
      "         -0.2312,  0.2979],\n",
      "        [-0.2382,  0.2633,  0.2655, -0.1762, -0.0289,  0.0482, -0.2789, -0.2819,\n",
      "          0.0546,  0.0115],\n",
      "        [ 0.0116,  0.0401,  0.2623, -0.1402,  0.1265, -0.2414, -0.0468,  0.1793,\n",
      "          0.0407,  0.0255],\n",
      "        [-0.1252, -0.0427,  0.1569, -0.1632, -0.1160,  0.2126,  0.0201, -0.0037,\n",
      "         -0.3134,  0.0699]], requires_grad=True), Parameter containing:\n",
      "tensor([-0.2379, -0.0660,  0.2595, -0.0033,  0.1164], requires_grad=True)]\n"
     ]
    }
   ],
   "source": [
    "class net1(nn.Module):\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.conv = nn.Linear(10,5)\n",
    "    self.tens = torch.ones(3,4) \n",
    "    \n",
    "  def forward(self, x):\n",
    "    return self.linear(x)\n",
    "\n",
    "myNet = net1()\n",
    "print(list(myNet.parameters()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5382e4d9-3e6b-423e-82b6-38a9d17e8fa3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Parameter containing:\n",
      "tensor([[1., 1., 1., 1.],\n",
      "        [1., 1., 1., 1.],\n",
      "        [1., 1., 1., 1.]], requires_grad=True), Parameter containing:\n",
      "tensor([[ 0.0675, -0.1424,  0.1650, -0.2677, -0.2785,  0.3050,  0.1304, -0.1756,\n",
      "          0.0919, -0.0243],\n",
      "        [-0.0186,  0.0041,  0.1209, -0.0411,  0.2620, -0.1601,  0.0917, -0.0100,\n",
      "          0.0036,  0.0561],\n",
      "        [-0.2973,  0.2082,  0.2299, -0.1645,  0.2716,  0.1492, -0.0174,  0.2474,\n",
      "         -0.3013, -0.2150],\n",
      "        [ 0.2291, -0.1482,  0.2357,  0.1155, -0.2492, -0.3007, -0.2570,  0.2553,\n",
      "         -0.2120, -0.0186],\n",
      "        [-0.2024, -0.3026, -0.0626, -0.2274,  0.0707, -0.0618,  0.1361, -0.0942,\n",
      "         -0.1200,  0.1006]], requires_grad=True), Parameter containing:\n",
      "tensor([-0.0016,  0.1711,  0.0386,  0.3026, -0.3049], requires_grad=True)]\n"
     ]
    }
   ],
   "source": [
    "class net2(nn.Module):\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.conv = nn.Linear(10,5) \n",
    "    self.tens = nn.Parameter(torch.ones(3,4))\n",
    "    \n",
    "  def forward(self, x):\n",
    "    return self.linear(x)\n",
    "\n",
    "\n",
    "myNet = net2()\n",
    "print(list(myNet.parameters()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1a38993d-11f8-4bf9-98b4-e7f80b962421",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Parameter containing:\n",
      "tensor([[ 0.1133, -0.1067, -0.0115, -0.2199,  0.0116,  0.1467, -0.0836,  0.0289,\n",
      "          0.1971,  0.1057],\n",
      "        [-0.0039, -0.1473,  0.2351,  0.0741,  0.1531, -0.2651,  0.0099, -0.0974,\n",
      "         -0.1724, -0.0945],\n",
      "        [ 0.1443, -0.1424, -0.2005,  0.0607, -0.0208,  0.2211, -0.1673, -0.1659,\n",
      "         -0.3016, -0.0901],\n",
      "        [-0.1222,  0.2375,  0.1785, -0.2714,  0.2844, -0.2399, -0.2231,  0.1849,\n",
      "          0.1270, -0.1058],\n",
      "        [-0.1723,  0.0931, -0.1755,  0.2425, -0.2066,  0.1930, -0.2187,  0.1941,\n",
      "          0.0350,  0.1038]], requires_grad=True), Parameter containing:\n",
      "tensor([ 0.1604, -0.0775, -0.1501, -0.2807, -0.0268], requires_grad=True), Parameter containing:\n",
      "tensor([[1., 1., 1., 1.],\n",
      "        [1., 1., 1., 1.],\n",
      "        [1., 1., 1., 1.]], requires_grad=True), Parameter containing:\n",
      "tensor([[ 0.2551,  0.2103, -0.2652, -0.2485,  0.2238, -0.2473,  0.0495, -0.1949,\n",
      "          0.1658, -0.2232],\n",
      "        [-0.0124, -0.1771, -0.2330, -0.2723, -0.2595,  0.1959, -0.1795, -0.0926,\n",
      "          0.1647, -0.1222],\n",
      "        [ 0.0052, -0.0189, -0.1721, -0.1111,  0.0831, -0.2996,  0.0761,  0.1099,\n",
      "         -0.0815, -0.0803],\n",
      "        [ 0.2251,  0.2358, -0.1614,  0.1889,  0.2141, -0.1088, -0.0604,  0.0556,\n",
      "          0.0349, -0.3086],\n",
      "        [ 0.1164,  0.2261, -0.0479, -0.2424, -0.2611,  0.2658,  0.1078, -0.0262,\n",
      "          0.0460,  0.2410]], requires_grad=True), Parameter containing:\n",
      "tensor([ 0.1287,  0.2776, -0.2042,  0.2004,  0.2015], requires_grad=True)]\n"
     ]
    }
   ],
   "source": [
    "class net3(nn.Module):\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.conv = nn.Linear(10,5) \n",
    "    self.net  = net2()\n",
    "    \n",
    "  def forward(self, x):\n",
    "    return self.linear(x)\n",
    "\n",
    "\n",
    "myNet = net3()\n",
    "print(list(myNet.parameters()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c5a68df-e468-403e-9954-537b2ae536f1",
   "metadata": {},
   "source": [
    "#### nn.ModuleList and nn.ParameterList"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "67b51e8c-51c2-478b-a86e-3126350f37af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    }
   ],
   "source": [
    "layer_list = [nn.Conv2d(5,5,3), nn.BatchNorm2d(5), nn.Linear(5,2)]\n",
    "\n",
    "class myNet(nn.Module):\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.layers = layer_list\n",
    "  \n",
    "  def forward(x):\n",
    "    for layer in self.layers:\n",
    "      x = layer(x)\n",
    "\n",
    "net = myNet()\n",
    "\n",
    "print(list(net.parameters())) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e707f30-c5f2-4465-a89b-7cabbaccb322",
   "metadata": {},
   "source": [
    "As you see, unlike when we would register individual modules, assigning a Python List doesn't register the parameters of Modules inside the list. To fix this, we wrap our list with the nn.ModuleList class, and then assign it as a member of the network class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ff2a40a8-a77c-400b-90d7-e792e129cbaf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Parameter containing:\n",
      "tensor([[[[-0.0248,  0.0262,  0.0701],\n",
      "          [ 0.0708,  0.1023, -0.0748],\n",
      "          [ 0.1310, -0.0304, -0.0038]],\n",
      "\n",
      "         [[ 0.1292, -0.0098, -0.0568],\n",
      "          [-0.0086,  0.1000,  0.0903],\n",
      "          [-0.0510, -0.1404, -0.0107]],\n",
      "\n",
      "         [[ 0.0700,  0.0678,  0.0759],\n",
      "          [-0.1122,  0.0283, -0.1265],\n",
      "          [-0.1018,  0.0413, -0.0480]],\n",
      "\n",
      "         [[ 0.0096,  0.1433,  0.0346],\n",
      "          [ 0.0484,  0.0780, -0.1206],\n",
      "          [-0.0040, -0.0600, -0.0282]],\n",
      "\n",
      "         [[ 0.0800,  0.1195,  0.0745],\n",
      "          [ 0.1180,  0.1145, -0.0248],\n",
      "          [ 0.0890, -0.1108,  0.0701]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0334, -0.1053, -0.0740],\n",
      "          [ 0.0668,  0.0229, -0.1259],\n",
      "          [ 0.0812, -0.1259, -0.0307]],\n",
      "\n",
      "         [[ 0.1404,  0.1040, -0.0406],\n",
      "          [-0.1438, -0.1406, -0.0927],\n",
      "          [-0.1313, -0.0680, -0.0595]],\n",
      "\n",
      "         [[-0.1411,  0.0346,  0.0881],\n",
      "          [ 0.0185,  0.0202,  0.0587],\n",
      "          [ 0.0161, -0.0022,  0.0624]],\n",
      "\n",
      "         [[ 0.0432, -0.0367, -0.1440],\n",
      "          [ 0.0580, -0.0080, -0.1174],\n",
      "          [ 0.0164, -0.0678, -0.0215]],\n",
      "\n",
      "         [[-0.1029,  0.0531,  0.0695],\n",
      "          [-0.1248, -0.0920,  0.0645],\n",
      "          [ 0.0359, -0.1189,  0.0760]]],\n",
      "\n",
      "\n",
      "        [[[-0.0237,  0.1185,  0.0501],\n",
      "          [ 0.0829, -0.1009,  0.0947],\n",
      "          [-0.0265, -0.0972, -0.0364]],\n",
      "\n",
      "         [[ 0.0011, -0.1409,  0.0667],\n",
      "          [ 0.0685,  0.1215,  0.1456],\n",
      "          [-0.1058,  0.0155, -0.1008]],\n",
      "\n",
      "         [[ 0.1319,  0.1194,  0.0534],\n",
      "          [ 0.1403, -0.0571,  0.1390],\n",
      "          [-0.0381, -0.0615,  0.0389]],\n",
      "\n",
      "         [[-0.0169, -0.0615,  0.0759],\n",
      "          [ 0.0944,  0.0314,  0.0432],\n",
      "          [-0.1238,  0.0541, -0.0534]],\n",
      "\n",
      "         [[ 0.0207,  0.1116,  0.0669],\n",
      "          [ 0.0673,  0.1146, -0.0624],\n",
      "          [-0.1164, -0.1370,  0.0650]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0348, -0.0415,  0.0065],\n",
      "          [-0.0098, -0.0674, -0.0953],\n",
      "          [-0.0397,  0.0385,  0.0039]],\n",
      "\n",
      "         [[ 0.0775, -0.0202, -0.0607],\n",
      "          [ 0.0326, -0.0964,  0.0601],\n",
      "          [ 0.1244,  0.0340,  0.1421]],\n",
      "\n",
      "         [[-0.0589, -0.0888, -0.1225],\n",
      "          [-0.0576,  0.0752, -0.0183],\n",
      "          [ 0.0735, -0.0740, -0.1250]],\n",
      "\n",
      "         [[-0.0349, -0.1091,  0.0633],\n",
      "          [-0.1255, -0.0476, -0.1088],\n",
      "          [-0.0067,  0.1477, -0.1485]],\n",
      "\n",
      "         [[-0.0273, -0.0712, -0.0679],\n",
      "          [ 0.0292,  0.0675,  0.0384],\n",
      "          [-0.0473, -0.0193, -0.0833]]],\n",
      "\n",
      "\n",
      "        [[[-0.0335,  0.1464, -0.1462],\n",
      "          [-0.0810, -0.0791, -0.1475],\n",
      "          [-0.0284,  0.1108,  0.1050]],\n",
      "\n",
      "         [[-0.1329, -0.0095,  0.0299],\n",
      "          [ 0.1119, -0.0590, -0.0497],\n",
      "          [-0.1374, -0.1234,  0.1351]],\n",
      "\n",
      "         [[ 0.0645,  0.0263, -0.0362],\n",
      "          [-0.0302,  0.1440, -0.1337],\n",
      "          [-0.0539, -0.0881,  0.1326]],\n",
      "\n",
      "         [[ 0.0099,  0.0101, -0.0417],\n",
      "          [ 0.0983, -0.1080,  0.0160],\n",
      "          [ 0.1022, -0.1198,  0.0556]],\n",
      "\n",
      "         [[-0.1476,  0.0970,  0.1200],\n",
      "          [ 0.0948,  0.1041,  0.0190],\n",
      "          [ 0.0267, -0.0053,  0.1295]]]], requires_grad=True), Parameter containing:\n",
      "tensor([-0.1011,  0.0808, -0.0836,  0.0901,  0.0017], requires_grad=True), Parameter containing:\n",
      "tensor([1., 1., 1., 1., 1.], requires_grad=True), Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0.], requires_grad=True), Parameter containing:\n",
      "tensor([[-0.4036, -0.1509,  0.2790,  0.4448, -0.2029],\n",
      "        [ 0.3466, -0.0573, -0.1602,  0.0857,  0.1873]], requires_grad=True), Parameter containing:\n",
      "tensor([0.1371, 0.1162], requires_grad=True)]\n"
     ]
    }
   ],
   "source": [
    "layer_list = [nn.Conv2d(5,5,3), nn.BatchNorm2d(5), nn.Linear(5,2)]\n",
    "\n",
    "class myNet(nn.Module):\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.layers = nn.ModuleList(layer_list)\n",
    "  \n",
    "  def forward(x):\n",
    "    for layer in self.layers:\n",
    "      x = layer(x)\n",
    "\n",
    "net = myNet()\n",
    "\n",
    "print(list(net.parameters())) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c71ac5f6-bb4a-4364-8a32-80ab1e02546e",
   "metadata": {},
   "source": [
    "Similarly, a list of tensors can be registered by wrapping the list inside a `nn.ParameterList` class.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd453ed1-642c-4dd3-b1ca-57ea092619cb",
   "metadata": {},
   "source": [
    "#### Weight Initialization\n",
    "\n",
    "Weight initialization plays a crucial role in the training of neural networks, as it can significantly influence the convergence speed and final performance of the model. Moreover, different types of layers may require different weight initialization schemes to ensure effective training.\n",
    "\n",
    "In PyTorch, the `modules()` function is a member function of the nn.Module class. It returns an iterator containing all the member nn.Module objects of the parent `nn.Module object. This allows us to access each individual layer within a neural network.\n",
    "\n",
    "Once we have access to each layer, we can use the apply() function to apply specific initialization schemes to the parameters of each layer. The apply() function takes as input a function or callable object that applies the desired initialization scheme to the parameters of the layer.\n",
    "\n",
    "Here's a simplified example of how we can use modules() and apply() to initialize the weights of different layers in a neural network:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "799b1428",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], requires_grad=True)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAigAAAGdCAYAAAA44ojeAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAaoklEQVR4nO3dcazV9X3/8dcF5IrKvfRauNcbQdCuVdfKDFa8rVu03gjIbE1pMxvSUUd0NRcTvVlbWVqZyxacMdPUUGmzDdZEZusWJNWUlmCFLb2g0pFWVkkxGFR6gUq4F+4v3ovc8/vjF89vt6L1IvR87uXxSL4J5/v9nHPf55sb7jPfe865dZVKpRIAgIKMqfUAAAC/TaAAAMURKABAcQQKAFAcgQIAFEegAADFESgAQHEECgBQnHG1HuBEDA4OZu/evZk4cWLq6upqPQ4A8B5UKpUcPnw4ra2tGTPm3a+RjMhA2bt3b6ZOnVrrMQCAE/DKK6/k/PPPf9c1IzJQJk6cmOT/PcGGhoYaTwMAvBe9vb2ZOnVq9ef4uxmRgfLWr3UaGhoECgCMMO/l5RleJAsAFEegAADFESgAQHEECgBQHIECABRHoAAAxREoAEBxBAoAUByBAgAUR6AAAMURKABAcQQKAFAcgQIAFEegAADFGVfrAYBTa/rdT9V6hGF7+b75tR4BqDFXUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4wwqU5cuX5+Mf/3gmTpyYKVOm5KabbsrOnTuHrHnjjTfS0dGRc889N+ecc04WLFiQffv2DVmzZ8+ezJ8/P2eddVamTJmSr3zlK3nzzTff/7MBAEaFYQXKpk2b0tHRkS1btmTDhg05evRorr/++vT19VXX3HXXXfnBD36Qxx9/PJs2bcrevXvz2c9+tnr82LFjmT9/fgYGBvLTn/40//qv/5rVq1fnnnvuOXnPCgAY0eoqlUrlRO984MCBTJkyJZs2bcqf/MmfpKenJ5MnT86aNWvyuc99Lkny4osv5pJLLklXV1euuuqq/PCHP8yf/umfZu/evWlubk6SrFy5Ml/72tdy4MCBjB8//nd+3d7e3jQ2NqanpycNDQ0nOj6cFqbf/VStRxi2l++bX+sRgFNgOD+/39drUHp6epIkTU1NSZJt27bl6NGjaW9vr665+OKLM23atHR1dSVJurq68rGPfawaJ0kyZ86c9Pb2ZseOHcf9Ov39/ent7R2yAQCj1wkHyuDgYO6888588pOfzEc/+tEkSXd3d8aPH59JkyYNWdvc3Jzu7u7qmv8dJ28df+vY8SxfvjyNjY3VberUqSc6NgAwApxwoHR0dOSFF17IY489djLnOa6lS5emp6enur3yyiun/GsCALUz7kTutGTJkjz55JPZvHlzzj///Or+lpaWDAwM5NChQ0Ououzbty8tLS3VNc8+++yQx3vrXT5vrflt9fX1qa+vP5FRAYARaFhXUCqVSpYsWZK1a9fm6aefzowZM4YcnzVrVs4444xs3Lixum/nzp3Zs2dP2trakiRtbW35xS9+kf3791fXbNiwIQ0NDbn00kvfz3MBAEaJYV1B6ejoyJo1a7Ju3bpMnDix+pqRxsbGTJgwIY2NjVm8eHE6OzvT1NSUhoaG3HHHHWlra8tVV12VJLn++utz6aWX5otf/GLuv//+dHd35+tf/3o6OjpcJQEAkgwzUB555JEkyTXXXDNk/6pVq/KlL30pSfLggw9mzJgxWbBgQfr7+zNnzpx861vfqq4dO3Zsnnzyydx+++1pa2vL2WefnUWLFuVv//Zv398zAQBGjff1OSi14nNQ4L3zOShAKX5vn4MCAHAqCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4oyr9QAAv2363U/VeoRhe/m++bUeAUYVV1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACjOuFoPACPJ9LufqvUIAKcFV1AAgOIIFACgOAIFACiOQAEAiiNQAIDiDDtQNm/enBtvvDGtra2pq6vLE088MeT4l770pdTV1Q3Z5s6dO2TNwYMHs3DhwjQ0NGTSpElZvHhxjhw58r6eCAAwegw7UPr6+jJz5sysWLHiHdfMnTs3v/71r6vbv/3bvw05vnDhwuzYsSMbNmzIk08+mc2bN+e2224b/vQAwKg07M9BmTdvXubNm/eua+rr69PS0nLcY7/85S+zfv36PPfcc7niiiuSJA8//HBuuOGGPPDAA2ltbR3uSADAKHNKXoPyzDPPZMqUKfnIRz6S22+/Pa+//nr1WFdXVyZNmlSNkyRpb2/PmDFjsnXr1uM+Xn9/f3p7e4dsAMDoddIDZe7cufnud7+bjRs35h/+4R+yadOmzJs3L8eOHUuSdHd3Z8qUKUPuM27cuDQ1NaW7u/u4j7l8+fI0NjZWt6lTp57ssQGAgpz0j7q/+eabq//+2Mc+lssuuywXXXRRnnnmmVx33XUn9JhLly5NZ2dn9XZvb69IAYBR7JS/zfjCCy/MBz/4wezatStJ0tLSkv379w9Z8+abb+bgwYPv+LqV+vr6NDQ0DNkAgNHrlAfKq6++mtdffz3nnXdekqStrS2HDh3Ktm3bqmuefvrpDA4OZvbs2ad6HABgBBj2r3iOHDlSvRqSJLt378727dvT1NSUpqam3HvvvVmwYEFaWlry0ksv5atf/Wo+9KEPZc6cOUmSSy65JHPnzs2tt96alStX5ujRo1myZEluvvlm7+ABAJKcwBWU559/Ppdffnkuv/zyJElnZ2cuv/zy3HPPPRk7dmx+/vOf59Of/nQ+/OEPZ/HixZk1a1b+8z//M/X19dXHePTRR3PxxRfnuuuuyw033JCrr7463/nOd07eswIARrRhX0G55pprUqlU3vH4j370o9/5GE1NTVmzZs1wvzQAcJrwt3gAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIozrtYDAIwG0+9+qtYjDNvL982v9QjwjlxBAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDjDDpTNmzfnxhtvTGtra+rq6vLEE08MOV6pVHLPPffkvPPOy4QJE9Le3p5f/epXQ9YcPHgwCxcuTENDQyZNmpTFixfnyJEj7+uJAACjx7ADpa+vLzNnzsyKFSuOe/z+++/PN7/5zaxcuTJbt27N2WefnTlz5uSNN96orlm4cGF27NiRDRs25Mknn8zmzZtz2223nfizAABGlXHDvcO8efMyb9684x6rVCp56KGH8vWvfz2f+cxnkiTf/e5309zcnCeeeCI333xzfvnLX2b9+vV57rnncsUVVyRJHn744dxwww154IEH0tra+j6eDgAwGpzU16Ds3r073d3daW9vr+5rbGzM7Nmz09XVlSTp6urKpEmTqnGSJO3t7RkzZky2bt163Mft7+9Pb2/vkA0AGL1OaqB0d3cnSZqbm4fsb25urh7r7u7OlClThhwfN25cmpqaqmt+2/Lly9PY2Fjdpk6dejLHBgAKMyLexbN06dL09PRUt1deeaXWIwEAp9BJDZSWlpYkyb59+4bs37dvX/VYS0tL9u/fP+T4m2++mYMHD1bX/Lb6+vo0NDQM2QCA0eukBsqMGTPS0tKSjRs3Vvf19vZm69ataWtrS5K0tbXl0KFD2bZtW3XN008/ncHBwcyePftkjgMAjFDDfhfPkSNHsmvXrurt3bt3Z/v27Wlqasq0adNy55135u/+7u/yB3/wB5kxY0a+8Y1vpLW1NTfddFOS5JJLLsncuXNz6623ZuXKlTl69GiWLFmSm2++2Tt4AIAkJxAozz//fK699trq7c7OziTJokWLsnr16nz1q19NX19fbrvtthw6dChXX3111q9fnzPPPLN6n0cffTRLlizJddddlzFjxmTBggX55je/eRKeDgAwGtRVKpVKrYcYrt7e3jQ2Nqanp8frUfi9mn73U7UeAU6al++bX+sROM0M5+f3iHgXDwBwehEoAEBxBAoAUByBAgAUR6AAAMURKABAcQQKAFAcgQIAFEegAADFESgAQHEECgBQHIECABRHoAAAxREoAEBxBAoAUByBAgAUR6AAAMURKABAcQQKAFAcgQIAFEegAADFESgAQHEECgBQHIECABRHoAAAxREoAEBxBAoAUByBAgAUR6AAAMURKABAcQQKAFCccbUegNPX9LufqvUIABTKFRQAoDgCBQAojkABAIojUACA4ggUAKA43sUDcJoaie+ke/m++bUegd8TV1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOKc9ED5m7/5m9TV1Q3ZLr744urxN954Ix0dHTn33HNzzjnnZMGCBdm3b9/JHgMAGMFOyRWUP/zDP8yvf/3r6vZf//Vf1WN33XVXfvCDH+Txxx/Ppk2bsnfv3nz2s589FWMAACPUuFPyoOPGpaWl5W37e3p68s///M9Zs2ZNPvWpTyVJVq1alUsuuSRbtmzJVVdddSrGAQBGmFNyBeVXv/pVWltbc+GFF2bhwoXZs2dPkmTbtm05evRo2tvbq2svvvjiTJs2LV1dXe/4eP39/ent7R2yAQCj10kPlNmzZ2f16tVZv359HnnkkezevTt//Md/nMOHD6e7uzvjx4/PpEmThtynubk53d3d7/iYy5cvT2NjY3WbOnXqyR4bACjISf8Vz7x586r/vuyyyzJ79uxccMEF+f73v58JEyac0GMuXbo0nZ2d1du9vb0iBQBGsVP+NuNJkyblwx/+cHbt2pWWlpYMDAzk0KFDQ9bs27fvuK9ZeUt9fX0aGhqGbADA6HXKA+XIkSN56aWXct5552XWrFk544wzsnHjxurxnTt3Zs+ePWlrazvVowAAI8RJ/xXPX/3VX+XGG2/MBRdckL1792bZsmUZO3ZsvvCFL6SxsTGLFy9OZ2dnmpqa0tDQkDvuuCNtbW3ewQMAVJ30QHn11VfzhS98Ia+//nomT56cq6++Olu2bMnkyZOTJA8++GDGjBmTBQsWpL+/P3PmzMm3vvWtkz0GADCC1VUqlUqthxiu3t7eNDY2pqenx+tRRrDpdz9V6xGAEebl++bXegTeh+H8/Pa3eACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOAIFACiOQAEAiiNQAIDinPSPuqc2fCorAKOJKygAQHEECgBQHIECABRHoAAAxREoAEBxBAoAUByBAgAUR6AAAMURKABAcQQKAFAcgQIAFEegAADFESgAQHEECgBQHIECABRHoAAAxREoAEBxBAoAUByBAgAUR6AAAMURKABAcQQKAFAcgQIAFEegAADFESgAQHEECgBQHIECABRHoAAAxREoAEBxBAoAUByBAgAUR6AAAMURKABAcQQKAFCccbUeAADeq+l3P1XrEYbt5fvm13qEEckVFACgOAIFACiOQAEAiiNQAIDiCBQAoDgCBQAojkABAIojUACA4ggUAKA4AgUAKI5AAQCKI1AAgOIIFACgOP6a8XGMxL+WCQCjiSsoAEBxBAoAUByBAgAUR6AAAMWpaaCsWLEi06dPz5lnnpnZs2fn2WefreU4AEAhahYo3/ve99LZ2Zlly5blZz/7WWbOnJk5c+Zk//79tRoJAChEzQLlH//xH3PrrbfmlltuyaWXXpqVK1fmrLPOyr/8y7/UaiQAoBA1+RyUgYGBbNu2LUuXLq3uGzNmTNrb29PV1fW29f39/env76/e7unpSZL09vaekvkG+//PKXlcAE4/0+56vNYjnJAX7p1z0h/zrZ/blUrld66tSaD85je/ybFjx9Lc3Dxkf3Nzc1588cW3rV++fHnuvffet+2fOnXqKZsRAE5njQ+dusc+fPhwGhsb33XNiPgk2aVLl6azs7N6e3BwMAcPHsy5556burq6d71vb29vpk6dmldeeSUNDQ2netQRz/kaPuds+Jyz4XG+hs85G77fxzmrVCo5fPhwWltbf+famgTKBz/4wYwdOzb79u0bsn/fvn1paWl52/r6+vrU19cP2Tdp0qRhfc2GhgbfpMPgfA2fczZ8ztnwOF/D55wN36k+Z7/ryslbavIi2fHjx2fWrFnZuHFjdd/g4GA2btyYtra2WowEABSkZr/i6ezszKJFi3LFFVfkyiuvzEMPPZS+vr7ccssttRoJAChEzQLlz/7sz3LgwIHcc8896e7uzh/90R9l/fr1b3vh7PtVX1+fZcuWve1XRByf8zV8ztnwOWfD43wNn3M2fKWds7rKe3mvDwDA75G/xQMAFEegAADFESgAQHEECgBQnNMqUD796U9n2rRpOfPMM3Peeefli1/8Yvbu3VvrsYr18ssvZ/HixZkxY0YmTJiQiy66KMuWLcvAwECtRyvW3//93+cTn/hEzjrrrGF/mODpYsWKFZk+fXrOPPPMzJ49O88++2ytRyrW5s2bc+ONN6a1tTV1dXV54oknaj1S8ZYvX56Pf/zjmThxYqZMmZKbbropO3furPVYxXrkkUdy2WWXVT+cra2tLT/84Q9rPVaS0yxQrr322nz/+9/Pzp078x//8R956aWX8rnPfa7WYxXrxRdfzODgYL797W9nx44defDBB7Ny5cr89V//da1HK9bAwEA+//nP5/bbb6/1KEX63ve+l87Ozixbtiw/+9nPMnPmzMyZMyf79++v9WhF6uvry8yZM7NixYpajzJibNq0KR0dHdmyZUs2bNiQo0eP5vrrr09fX1+tRyvS+eefn/vuuy/btm3L888/n0996lP5zGc+kx07dtR6tKRyGlu3bl2lrq6uMjAwUOtRRoz777+/MmPGjFqPUbxVq1ZVGhsbaz1Gca688spKR0dH9faxY8cqra2tleXLl9dwqpEhSWXt2rW1HmPE2b9/fyVJZdOmTbUeZcT4wAc+UPmnf/qnWo9ROa2uoPxvBw8ezKOPPppPfOITOeOMM2o9zojR09OTpqamWo/BCDQwMJBt27alvb29um/MmDFpb29PV1dXDSdjNOvp6UkS/2+9B8eOHctjjz2Wvr6+Iv7szGkXKF/72tdy9tln59xzz82ePXuybt26Wo80YuzatSsPP/xw/vIv/7LWozAC/eY3v8mxY8fe9mnRzc3N6e7urtFUjGaDg4O5884788lPfjIf/ehHaz1OsX7xi1/knHPOSX19fb785S9n7dq1ufTSS2s91sgPlLvvvjt1dXXvur344ovV9V/5ylfy3//93/nxj3+csWPH5s///M9TOc0+THe45yxJXnvttcydOzef//znc+utt9Zo8to4kfMF1F5HR0deeOGFPPbYY7UepWgf+chHsn379mzdujW33357Fi1alP/5n/+p9Vgj/6PuDxw4kNdff/1d11x44YUZP3782/a/+uqrmTp1an76058WcTnr92W452zv3r255pprctVVV2X16tUZM2bEd+2wnMj32OrVq3PnnXfm0KFDp3i6kWNgYCBnnXVW/v3f/z033XRTdf+iRYty6NAhVzN/h7q6uqxdu3bIueOdLVmyJOvWrcvmzZszY8aMWo8zorS3t+eiiy7Kt7/97ZrOUbM/FniyTJ48OZMnTz6h+w4ODiZJ+vv7T+ZIxRvOOXvttddy7bXXZtasWVm1atVpFyfJ+/se4/8bP358Zs2alY0bN1Z/yA4ODmbjxo1ZsmRJbYdj1KhUKrnjjjuydu3aPPPMM+LkBAwODhbxc3HEB8p7tXXr1jz33HO5+uqr84EPfCAvvfRSvvGNb+Siiy46ra6eDMdrr72Wa665JhdccEEeeOCBHDhwoHqspaWlhpOVa8+ePTl48GD27NmTY8eOZfv27UmSD33oQznnnHNqO1wBOjs7s2jRolxxxRW58sor89BDD6Wvry+33HJLrUcr0pEjR7Jr167q7d27d2f79u1pamrKtGnTajhZuTo6OrJmzZqsW7cuEydOrL6+qbGxMRMmTKjxdOVZunRp5s2bl2nTpuXw4cNZs2ZNnnnmmfzoRz+q9Winz9uMf/7zn1euvfbaSlNTU6W+vr4yffr0ype//OXKq6++WuvRirVq1apKkuNuHN+iRYuOe75+8pOf1Hq0Yjz88MOVadOmVcaPH1+58sorK1u2bKn1SMX6yU9+ctzvp0WLFtV6tGK90/9Zq1atqvVoRfqLv/iLygUXXFAZP358ZfLkyZXrrruu8uMf/7jWY1UqlUplxL8GBQAYfU6/FxQAAMUTKABAcQQKAFAcgQIAFEegAADFESgAQHEECgBQHIECABRHoAAAxREoAEBxBAoAUByBAgAU5/8CtVe+6o/KLvEAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "class myNet(nn.Module):\n",
    " \n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.conv = nn.Conv2d(10,10,3)\n",
    "    self.bn = nn.BatchNorm2d(10)\n",
    "  \n",
    "  def weights_init(self):\n",
    "    for module in self.modules():\n",
    "        if isinstance(module, nn.Conv2d):\n",
    "            nn.init.normal_(module.weight, mean = 0, std = 1)\n",
    "            nn.init.constant_(module.bias, 0)\n",
    "\n",
    "Net = myNet()\n",
    "Net.weights_init()\n",
    "\n",
    "for module in Net.modules():\n",
    "    if isinstance(module, nn.Conv2d):\n",
    "        weights = module.weight\n",
    "        weights = weights.reshape(-1).detach().cpu().numpy()\n",
    "        print(module.bias)\n",
    "        plt.hist(weights)\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb70395b",
   "metadata": {},
   "source": [
    "#### modules vs children\n",
    "\n",
    "`modules()` and `children()` are both functions in PyTorch's `nn.Module` class that allow us to access the submodules (i.e., layers) within a neural network. While they serve similar purposes, there is a subtle but important difference between them.\n",
    "\n",
    "The `children()` function returns an iterable containing only the direct child modules (i.e., immediate submodules) of the parent module. These child modules are the ones directly assigned as attributes of the parent module.\n",
    "\n",
    "On the other hand, the `modules()` function recursively goes through all submodules of the parent module, including submodules of submodules, and so on. It returns an iterable containing all modules within the hierarchy of the parent module.\n",
    "\n",
    "Here's a simple example to illustrate the difference between `children()` and `modules()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "533696e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Printing children\n",
      "------------------------------\n",
      "[Sequential(\n",
      "  (0): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
      "  (1): BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "), Linear(in_features=10, out_features=2, bias=True)]\n",
      "\n",
      "\n",
      "Printing Modules\n",
      "------------------------------\n",
      "[MyNet(\n",
      "  (convBN): Sequential(\n",
      "    (0): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
      "    (1): BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (linear): Linear(in_features=10, out_features=2, bias=True)\n",
      "), Sequential(\n",
      "  (0): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
      "  (1): BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "), Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1)), BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), Linear(in_features=10, out_features=2, bias=True)]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class MyNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.convBN = nn.Sequential(nn.Conv2d(10, 10, 3), nn.BatchNorm2d(10))\n",
    "        self.linear = nn.Linear(10, 2)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        pass\n",
    "\n",
    "net = MyNet()\n",
    "\n",
    "print(\"Printing children\\n------------------------------\")\n",
    "print(list(net.children()))\n",
    "print(\"\\n\\nPrinting Modules\\n------------------------------\")\n",
    "print(list(net.modules()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "495f7032",
   "metadata": {},
   "source": [
    "#### Printing Information About the Network\n",
    "\n",
    "Printing information about a neural network in PyTorch can be useful for both users and debugging purposes. PyTorch provides several named functions to access information about the network's parameters, modules, children, and buffers. These functions allow us to iterate through the network's components while also providing their names for easy identification.\n",
    "\n",
    "Here's a brief overview of the four named functions available in PyTorch:\n",
    "\n",
    "1. `named_parameters`: Returns an iterator containing tuples of parameter names and their corresponding values. Parameters typically include weights and biases of the network's layers.\n",
    "\n",
    "2. `named_modules`: Similar to named_parameters, but returns an iterator containing tuples of module names and their corresponding modules. This function traverses through all modules within the network, including nested modules.\n",
    "\n",
    "3. `named_children`: Like named_modules, but specifically returns an iterator containing tuples of child module names and their corresponding modules. This function only traverses through direct child modules of the network.\n",
    "\n",
    "4. `named_buffers`: Returns an iterator containing tuples of buffer names and their corresponding tensor values. Buffers often include non-learnable parameters, such as running mean and variance in batch normalization layers.\n",
    "\n",
    "Here's a simple example demonstrating the usage of these named functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6c0f0802",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Named Parameters:\n",
      "conv1.weight torch.Size([64, 3, 3, 3])\n",
      "conv1.bias torch.Size([64])\n",
      "bn1.weight torch.Size([64])\n",
      "bn1.bias torch.Size([64])\n",
      "fc.weight torch.Size([10, 50176])\n",
      "fc.bias torch.Size([10])\n",
      "\n",
      "Named Modules:\n",
      " MyNet(\n",
      "  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (relu): ReLU()\n",
      "  (fc): Linear(in_features=50176, out_features=10, bias=True)\n",
      ")\n",
      "conv1 Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "bn1 BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "relu ReLU()\n",
      "fc Linear(in_features=50176, out_features=10, bias=True)\n",
      "\n",
      "Named Children:\n",
      "conv1 Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "bn1 BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "relu ReLU()\n",
      "fc Linear(in_features=50176, out_features=10, bias=True)\n",
      "\n",
      "Named Buffers:\n",
      "bn1.running_mean torch.Size([64])\n",
      "bn1.running_var torch.Size([64])\n",
      "bn1.num_batches_tracked torch.Size([])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# Define a custom neural network\n",
    "class MyNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(64)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.fc = nn.Linear(64 * 28 * 28, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        pass\n",
    "\n",
    "# Create an instance of the network\n",
    "net = MyNet()\n",
    "\n",
    "# Print information about the network's parameters\n",
    "print(\"Named Parameters:\")\n",
    "for name, param in net.named_parameters():\n",
    "    print(name, param.shape)\n",
    "\n",
    "# Print information about the network's modules\n",
    "print(\"\\nNamed Modules:\")\n",
    "for name, module in net.named_modules():\n",
    "    print(name, module)\n",
    "\n",
    "# Print information about the network's children\n",
    "print(\"\\nNamed Children:\")\n",
    "for name, child in net.named_children():\n",
    "    print(name, child)\n",
    "\n",
    "# Print information about the network's buffers\n",
    "print(\"\\nNamed Buffers:\")\n",
    "for name, buffer in net.named_buffers():\n",
    "    print(name, buffer.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2d42256",
   "metadata": {},
   "source": [
    "#### Differential Learning Rate\n",
    "\n",
    "This capability allows us to fine-tune the training process by adjusting the learning rates based on the specific requirements of different parts of the network.\n",
    "\n",
    "The implementation of this concept is straightforward. In our previous example where we created a CIFAR classifier, we passed all the network's parameters as a single group to the optimizer object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d1016168",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(10, 5)\n",
    "        self.fc2 = nn.Linear(5, 2)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.fc2(self.fc1(x))\n",
    "\n",
    "net = MyNet()\n",
    "optimizer = torch.optim.SGD(net.parameters(), lr=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90877340",
   "metadata": {},
   "source": [
    "However, PyTorch's `torch.optim` class allows us to specify different sets of parameters with distinct learning rates using a dictionary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1e08eda6",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.SGD([\n",
    "    {\"params\": net.fc1.parameters(), 'lr': 0.001, \"momentum\": 0.99},\n",
    "    {\"params\": net.fc2.parameters()}\n",
    "], lr=0.01, momentum=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a9d28ca",
   "metadata": {},
   "source": [
    "In this scenario, the parameters of `fc1` are assigned a learning rate of `0.001` and a momentum of `0.99`. If a hyperparameter is not specified for a group of parameters (like `fc2`), they use the default value of that hyperparameter provided as an input argument to the optimizer function.\n",
    "\n",
    "Alternatively, if we want to set different hyperparameters for biases and weights separately, we can do so by separating them and then configuring the optimizer accordingly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1464d411",
   "metadata": {},
   "outputs": [],
   "source": [
    "params_bias = []\n",
    "params_weights = []\n",
    "\n",
    "# Separate the bias and weight parameters\n",
    "for name, parameter in net.named_parameters():\n",
    "    if \"bias\" in name:\n",
    "        params_bias.append(parameter)\n",
    "    elif \"weight\" in name:\n",
    "        params_weights.append(parameter)\n",
    "\n",
    "# Set the optimizer to have different hyperparameters for biases and weights\n",
    "optimizer = torch.optim.SGD([\n",
    "    {\"params\": params_bias, 'lr': 0.001, \"momentum\": 0.99},\n",
    "    {\"params\": params_weights}\n",
    "], lr=0.01, momentum=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87db7054",
   "metadata": {},
   "source": [
    "This approach allows us to fine-tune the training process by customizing the learning rates for different layers or types of parameters in the neural network."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ad0ab15",
   "metadata": {},
   "source": [
    "#### Scheduling Learning Rates\n",
    "\n",
    "Scheduling learning rates is a critical aspect of training neural networks, and PyTorch offers support for various learning rate schedules through its torch.optim.lr_scheduler module. One such example is demonstrated below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b3605ab7",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[10, 20], gamma=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fb3165f",
   "metadata": {},
   "source": [
    "In the above example, the scheduler multiplies the learning rate by the specified `gamma` factor each time the training reaches an epoch contained in the `milestones` list. For instance, the learning rate is multiplied by 0.1 at the 10th and 20th epochs.\n",
    "\n",
    "To integrate this scheduler into your training loop, you need to ensure that you call `scheduler.step()` at the beginning of the epoch loop. Typically, the training loop consists of two nested loops: one iterating over epochs and the other over batches within each epoch. The `scheduler.step()` should be placed outside the batch loop to update the learning rate at the start of each epoch. Avoid placing it inside the batch loop, as this may lead to unintended updates to the learning rate.\n",
    "\n",
    "It's important to note that `scheduler.step()` should not replace `optimizer.step()`. You still need to call `optimizer.step()` every time you perform backpropagation (inside the batch loop), while `scheduler.step()` updates the learning rate at the beginning of each epoch.\n",
    "\n",
    "Here's a basic structure of how you would integrate the scheduler into your training loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6254ab74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[10, 20], gamma=0.1)\n",
    "\n",
    "# for epoch in range(num_epochs):\n",
    "#     scheduler.step()  # Update learning rate at the start of each epoch\n",
    "    \n",
    "#     # Training loop\n",
    "#     for batch_idx, (inputs, targets) in enumerate(train_loader):\n",
    "#         optimizer.zero_grad()\n",
    "#         outputs = model(inputs)\n",
    "#         loss = criterion(outputs, targets)\n",
    "#         loss.backward()\n",
    "#         optimizer.step()  # Backpropagation\n",
    "        \n",
    "#     # Validation loop (if applicable)\n",
    "#     # Calculate validation loss and other metrics\n",
    "    \n",
    "#     print(f\"Epoch {epoch+1}/{num_epochs}, Learning Rate: {scheduler.get_last_lr()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dbb7ad6",
   "metadata": {},
   "source": [
    "In this example, scheduler.get_last_lr() retrieves the learning rate after the scheduler has been updated. Adjust this structure according to your specific training setup and requirements."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "304f622b",
   "metadata": {},
   "source": [
    "#### Saving Model\n",
    "\n",
    "Saving your model in PyTorch can be achieved in two ways: using `torch.save()` to save the entire model or using the model's `state_dict` to save only the parameters. Let's explore both methods:\n",
    "\n",
    "**Saving the Entire Model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f5480a64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "myNet(\n",
      "  (conv): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
      "  (bn): BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Save the entire model\n",
    "torch.save(Net, \"net.pth\")\n",
    "\n",
    "# Load the model\n",
    "Net = torch.load(\"net.pth\")\n",
    "\n",
    "# Print the loaded model\n",
    "print(Net)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a59269a",
   "metadata": {},
   "source": [
    "In this approach, the entire model, including its architecture and weights, is serialized and saved to disk. Later, you can load the model back into memory using `torch.load()`.\n",
    "\n",
    "**Saving and Loading Model's State Dict**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b6213057",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "conv.weight tensor([[[[-1.0087e+00, -1.3472e-01,  5.5503e-01],\n",
      "          [ 7.5895e-01,  2.0397e-01,  6.3517e-01],\n",
      "          [-6.7221e-01, -2.3866e+00,  2.3710e-01]],\n",
      "\n",
      "         [[-2.1001e+00, -4.5312e-01, -1.2131e+00],\n",
      "          [ 4.4503e-02, -2.4762e-01,  1.5118e+00],\n",
      "          [ 1.0428e+00,  3.0670e-01,  1.8215e+00]],\n",
      "\n",
      "         [[ 4.9656e-01,  1.0795e+00,  1.2628e+00],\n",
      "          [ 5.6638e-02,  6.2805e-01, -9.7419e-01],\n",
      "          [-1.1829e+00, -4.9368e-01,  1.2095e+00]],\n",
      "\n",
      "         [[-2.1280e-01, -1.4508e+00, -1.2144e+00],\n",
      "          [ 1.3363e+00, -6.7043e-01, -2.0833e-01],\n",
      "          [-3.5866e-01,  1.2548e+00,  1.1278e+00]],\n",
      "\n",
      "         [[ 5.1350e-01, -9.2347e-01,  3.0960e-01],\n",
      "          [-4.4869e-02,  1.2497e+00,  9.3222e-01],\n",
      "          [-1.9099e-01, -4.5086e-01,  8.8910e-01]],\n",
      "\n",
      "         [[-1.7785e+00, -5.9751e-01, -2.7587e-01],\n",
      "          [ 1.8195e-01, -1.0054e+00, -2.6742e-01],\n",
      "          [-5.4385e-01,  2.8419e-01,  8.2707e-02]],\n",
      "\n",
      "         [[ 1.0776e+00, -9.5793e-01,  5.3833e-01],\n",
      "          [-4.4167e-01, -1.9093e-02, -4.8108e-01],\n",
      "          [-3.5484e-01,  3.8840e-01,  6.5141e-01]],\n",
      "\n",
      "         [[-4.7164e-01,  1.8347e+00, -1.6562e+00],\n",
      "          [ 1.3879e+00,  6.1261e-01, -4.1332e-01],\n",
      "          [-8.9406e-01, -7.0546e-02, -2.4055e+00]],\n",
      "\n",
      "         [[ 1.2746e+00, -1.5707e-02, -4.7210e-01],\n",
      "          [-7.2779e-01, -5.2791e-02, -4.4878e-01],\n",
      "          [-5.3615e-01,  1.3339e+00,  1.3059e+00]],\n",
      "\n",
      "         [[-1.2561e-01, -7.3947e-01, -8.7954e-01],\n",
      "          [-1.5197e+00, -6.0248e-01, -5.9761e-01],\n",
      "          [-3.5896e-01,  6.5164e-01, -5.0170e-01]]],\n",
      "\n",
      "\n",
      "        [[[-1.4819e+00,  1.5623e-01,  9.7013e-01],\n",
      "          [ 1.3281e+00,  4.8476e-01,  2.1240e-01],\n",
      "          [ 1.1274e+00, -6.2731e-01,  7.2656e-01]],\n",
      "\n",
      "         [[ 9.1362e-01,  4.4857e-01, -1.0254e+00],\n",
      "          [-5.0960e-01, -5.4690e-01, -1.3217e-01],\n",
      "          [-1.1158e+00, -4.6894e-01,  1.3054e+00]],\n",
      "\n",
      "         [[-6.1155e-01,  8.0145e-01, -4.2879e-01],\n",
      "          [ 1.0923e+00, -2.0537e-02, -1.0006e+00],\n",
      "          [ 9.8836e-02, -4.1458e-01, -4.9854e-02]],\n",
      "\n",
      "         [[ 1.7741e-01, -8.5228e-01,  5.8873e-01],\n",
      "          [ 9.5382e-01, -1.1157e+00, -1.7948e+00],\n",
      "          [-7.6153e-01, -3.0033e-01,  6.8677e-01]],\n",
      "\n",
      "         [[-1.3953e+00,  3.6441e-01,  1.1982e+00],\n",
      "          [ 1.2281e+00, -6.3720e-01,  3.6823e-01],\n",
      "          [-7.9624e-01,  1.5863e+00, -5.4826e-01]],\n",
      "\n",
      "         [[-1.9080e-01,  1.0512e+00, -8.9455e-02],\n",
      "          [-2.3663e+00, -7.7120e-01, -4.9867e-01],\n",
      "          [-5.4372e-01, -1.6064e+00, -7.3879e-01]],\n",
      "\n",
      "         [[-9.4028e-01, -4.2245e-01,  7.8466e-02],\n",
      "          [ 6.8112e-01, -6.9630e-02, -8.7597e-01],\n",
      "          [-2.0574e+00, -1.1639e+00,  1.3891e+00]],\n",
      "\n",
      "         [[ 4.8526e-01,  3.1257e-01,  4.1627e-01],\n",
      "          [ 4.2019e-01, -4.7033e-01, -7.5785e-01],\n",
      "          [-7.6427e-01,  4.6424e-01, -1.1510e+00]],\n",
      "\n",
      "         [[-1.3529e+00,  1.1982e-01, -1.3723e+00],\n",
      "          [ 1.2337e+00,  8.8937e-01, -4.5746e-01],\n",
      "          [-5.5449e-01,  1.6034e+00, -6.5170e-01]],\n",
      "\n",
      "         [[ 6.8308e-01, -2.1594e-01,  1.6004e+00],\n",
      "          [ 6.7027e-01,  4.1551e-01, -1.4082e+00],\n",
      "          [-8.0225e-01, -1.3048e+00,  1.2917e+00]]],\n",
      "\n",
      "\n",
      "        [[[ 3.2148e-01,  2.5177e-01,  1.4496e+00],\n",
      "          [ 5.4600e-01,  4.2928e-01,  1.9329e+00],\n",
      "          [ 1.5977e+00, -6.1818e-04,  1.0415e+00]],\n",
      "\n",
      "         [[-1.2549e+00,  2.8533e-01, -3.1235e-01],\n",
      "          [-3.3967e-01, -8.3695e-01, -1.1368e-01],\n",
      "          [ 5.5915e-02,  1.4461e+00, -1.3086e-01]],\n",
      "\n",
      "         [[ 8.6650e-01, -3.6607e-02, -5.9474e-03],\n",
      "          [-3.2974e-01, -8.3209e-01, -3.2175e-01],\n",
      "          [-2.1081e-01,  1.1055e-01,  8.6922e-01]],\n",
      "\n",
      "         [[ 2.0795e+00, -7.7919e-01,  4.7581e-01],\n",
      "          [ 1.8334e+00, -2.8864e-01, -4.0183e-01],\n",
      "          [ 1.0266e+00,  1.7017e-01,  1.7088e+00]],\n",
      "\n",
      "         [[-6.1966e-01,  2.6876e-01, -2.8465e+00],\n",
      "          [ 9.9266e-01, -1.1801e+00, -5.1291e-02],\n",
      "          [ 4.2673e-01,  1.4978e-01,  7.8353e-03]],\n",
      "\n",
      "         [[-1.8885e+00, -1.7163e-01, -1.1282e+00],\n",
      "          [ 4.9594e-01, -1.0436e+00, -6.1740e-01],\n",
      "          [-7.2360e-01,  1.8834e-01, -3.0624e-02]],\n",
      "\n",
      "         [[ 1.5044e+00, -1.7102e+00, -1.9355e-01],\n",
      "          [-1.4625e+00, -2.3314e-01,  4.1163e-01],\n",
      "          [-1.0870e+00, -6.1593e-02,  2.2074e-01]],\n",
      "\n",
      "         [[ 1.6812e+00, -1.9894e+00,  4.8241e-02],\n",
      "          [ 4.7644e-01, -7.2524e-01,  1.8385e+00],\n",
      "          [-5.1645e-01, -3.6038e-01, -7.2618e-01]],\n",
      "\n",
      "         [[-9.0583e-01, -1.3234e+00, -3.7648e-01],\n",
      "          [-1.0948e-01, -1.7913e+00, -1.4616e+00],\n",
      "          [ 2.6488e-01,  8.7914e-01, -8.7993e-01]],\n",
      "\n",
      "         [[-1.8167e+00,  1.3124e+00,  7.9601e-01],\n",
      "          [-1.0830e+00,  1.1401e+00, -8.2844e-01],\n",
      "          [ 4.0597e-01,  1.0469e-02,  5.7592e-02]]],\n",
      "\n",
      "\n",
      "        [[[-1.3331e-01, -2.3507e-01, -9.3880e-01],\n",
      "          [ 1.5502e+00, -2.6373e-01,  7.9435e-01],\n",
      "          [ 1.5033e-01,  3.5882e-01,  1.3200e-01]],\n",
      "\n",
      "         [[-3.1506e-01, -9.8456e-02,  4.3422e-02],\n",
      "          [-1.3070e-01, -6.3103e-01, -1.6607e+00],\n",
      "          [ 5.7313e-01,  7.3234e-01, -1.1117e+00]],\n",
      "\n",
      "         [[-4.4557e-02,  1.5643e+00, -1.6169e+00],\n",
      "          [-7.3236e-01,  5.3045e-01, -8.4651e-01],\n",
      "          [-1.6458e-01,  1.0941e+00, -7.8636e-02]],\n",
      "\n",
      "         [[ 2.4047e-01, -3.1207e-01,  6.8154e-01],\n",
      "          [-4.5610e-01,  2.8918e-01, -9.2402e-01],\n",
      "          [-6.7271e-01,  8.5481e-01,  1.4954e+00]],\n",
      "\n",
      "         [[-9.8192e-02,  1.6234e+00,  8.9171e-01],\n",
      "          [-1.0001e+00,  1.7078e-02, -7.7003e-01],\n",
      "          [-5.4652e-01, -6.7070e-02,  1.5807e+00]],\n",
      "\n",
      "         [[ 1.1380e-01,  5.7846e-01,  1.8611e-01],\n",
      "          [-4.0756e-01,  2.6428e-01, -1.7906e-01],\n",
      "          [-9.7697e-01,  9.6907e-01, -5.6938e-01]],\n",
      "\n",
      "         [[-1.0041e+00,  4.6178e-01,  1.0493e+00],\n",
      "          [-3.0588e-01, -5.1625e-01,  1.2141e+00],\n",
      "          [ 1.7869e+00, -4.4820e-01, -9.9904e-01]],\n",
      "\n",
      "         [[ 8.0144e-01, -8.0010e-01, -6.0493e-01],\n",
      "          [-3.4150e-01, -1.8933e+00, -1.7255e+00],\n",
      "          [ 1.0960e+00,  8.7046e-01, -1.1102e+00]],\n",
      "\n",
      "         [[-5.8367e-01,  1.8273e-01, -1.0747e-01],\n",
      "          [-4.5553e-02,  7.5472e-01,  1.0404e+00],\n",
      "          [-1.9795e-02, -7.6912e-01, -1.0775e+00]],\n",
      "\n",
      "         [[ 2.3733e-01,  2.1361e+00, -5.1236e-01],\n",
      "          [-1.2854e+00, -1.4740e+00, -4.2206e-01],\n",
      "          [ 4.1466e-02,  2.8522e-01, -6.1284e-02]]],\n",
      "\n",
      "\n",
      "        [[[-7.6109e-01,  9.6878e-01,  4.1644e-01],\n",
      "          [-1.0822e+00,  5.4170e-03,  5.6751e-01],\n",
      "          [ 8.7653e-01, -7.9907e-01,  1.9674e+00]],\n",
      "\n",
      "         [[-7.0771e-01, -4.9960e-01, -2.9193e-01],\n",
      "          [-2.3709e-01, -6.1371e-02,  5.2126e-01],\n",
      "          [ 4.6963e-02, -6.6984e-01,  1.6528e-02]],\n",
      "\n",
      "         [[ 2.6758e-01,  1.2969e-02,  5.9238e-01],\n",
      "          [-1.1635e+00,  7.3007e-01, -9.1967e-01],\n",
      "          [-1.2998e-01, -4.6051e-01, -1.2131e+00]],\n",
      "\n",
      "         [[-3.8354e-03, -9.6301e-01,  7.2848e-01],\n",
      "          [-8.1794e-01, -1.6001e+00, -1.1064e+00],\n",
      "          [-1.6282e+00, -3.1247e-01,  1.7612e+00]],\n",
      "\n",
      "         [[ 7.5443e-02,  8.1512e-01,  5.1581e-02],\n",
      "          [ 9.6755e-02,  7.7988e-01, -5.2042e-01],\n",
      "          [-1.6369e+00,  8.2963e-02,  6.8710e-01]],\n",
      "\n",
      "         [[ 2.3165e-01, -7.5956e-01, -1.1646e+00],\n",
      "          [-8.5268e-01,  1.5265e+00, -9.1271e-01],\n",
      "          [-1.0869e+00,  5.6429e-01, -1.7038e+00]],\n",
      "\n",
      "         [[ 2.7756e-01, -1.6735e+00, -1.9490e-01],\n",
      "          [-7.7180e-01, -5.6471e-01,  8.8456e-01],\n",
      "          [-5.2857e-01, -1.0247e-01, -1.1878e+00]],\n",
      "\n",
      "         [[-4.1120e-02,  5.6541e-01, -2.0265e-01],\n",
      "          [-1.2630e+00,  1.1587e+00, -6.1204e-01],\n",
      "          [ 1.1095e+00, -1.3783e+00,  1.2844e+00]],\n",
      "\n",
      "         [[ 6.5416e-02,  1.4566e+00,  1.1315e+00],\n",
      "          [ 1.9247e+00, -1.2537e-01,  1.3419e-01],\n",
      "          [-5.2269e-01,  1.3386e-02,  6.1552e-02]],\n",
      "\n",
      "         [[-1.3556e+00, -3.6380e-01,  1.4651e+00],\n",
      "          [ 4.6071e-01,  4.0575e-01,  4.4783e-01],\n",
      "          [ 9.3964e-01, -1.4772e+00,  1.3444e-01]]],\n",
      "\n",
      "\n",
      "        [[[ 4.5488e-01, -3.7599e-01, -5.1097e-02],\n",
      "          [-5.4233e-01,  1.1490e+00, -2.6164e-02],\n",
      "          [ 1.1379e-01,  1.8825e-01,  7.2691e-01]],\n",
      "\n",
      "         [[ 7.3633e-01, -7.8313e-01, -1.8414e+00],\n",
      "          [ 1.0232e+00,  2.1491e-01,  8.8228e-01],\n",
      "          [ 5.5072e-01,  5.8687e-02,  1.4150e+00]],\n",
      "\n",
      "         [[ 1.3103e+00,  6.3364e-01,  3.6840e-01],\n",
      "          [-2.2025e+00, -1.9180e+00, -2.0189e-01],\n",
      "          [-7.5986e-02, -6.9578e-01,  2.8902e-01]],\n",
      "\n",
      "         [[-9.4163e-01, -1.2281e+00,  2.1820e-02],\n",
      "          [ 3.7465e-01, -4.9254e-01,  1.5797e+00],\n",
      "          [ 1.8777e+00,  3.4879e-01,  1.1768e+00]],\n",
      "\n",
      "         [[ 1.3092e+00,  9.6060e-01,  1.2818e+00],\n",
      "          [-1.9008e+00, -1.6186e+00, -7.2717e-01],\n",
      "          [-1.1510e+00, -3.6345e-01,  1.6259e+00]],\n",
      "\n",
      "         [[ 4.7612e-01,  1.1622e+00, -1.6697e+00],\n",
      "          [-6.7895e-01, -1.7926e+00, -4.6558e-01],\n",
      "          [ 1.7558e+00,  2.1208e-01, -4.7317e-02]],\n",
      "\n",
      "         [[ 1.7782e+00, -2.6806e-01, -7.2287e-02],\n",
      "          [ 8.4139e-01,  4.4672e-01, -4.3376e-01],\n",
      "          [ 1.4188e+00, -3.0883e-01, -4.5428e-02]],\n",
      "\n",
      "         [[ 3.5926e-02, -2.6033e-01, -5.4734e-02],\n",
      "          [ 3.7009e-02,  1.0014e+00, -5.2190e-01],\n",
      "          [ 6.2889e-01, -1.2496e+00,  4.7444e-01]],\n",
      "\n",
      "         [[-7.4718e-02, -6.3140e-01,  8.9580e-01],\n",
      "          [ 1.3104e+00, -2.1098e+00, -2.4088e+00],\n",
      "          [ 3.2981e-01, -1.1288e-01, -1.4148e+00]],\n",
      "\n",
      "         [[-1.2010e+00,  5.3013e-01, -7.2937e-01],\n",
      "          [ 1.2014e+00, -1.8197e-01, -3.9908e-01],\n",
      "          [ 5.6609e-02,  4.5417e-01,  5.0973e-01]]],\n",
      "\n",
      "\n",
      "        [[[ 6.3006e-02, -5.0553e-01, -2.1294e+00],\n",
      "          [ 5.3779e-01, -1.2969e+00,  8.8071e-01],\n",
      "          [-3.8384e-01, -1.2045e+00, -9.6792e-01]],\n",
      "\n",
      "         [[ 2.6610e-01,  5.1668e-01, -1.3500e-01],\n",
      "          [-1.5274e-02,  9.0745e-01,  7.7541e-01],\n",
      "          [ 1.1717e+00, -7.8266e-01,  8.9805e-01]],\n",
      "\n",
      "         [[ 1.4131e+00,  3.0987e-02, -1.4720e+00],\n",
      "          [ 1.3699e+00,  1.7859e+00, -1.3532e-01],\n",
      "          [-1.7951e-01, -8.5745e-01,  1.3581e+00]],\n",
      "\n",
      "         [[-1.0510e+00, -3.2023e-01,  5.3305e-01],\n",
      "          [-9.1891e-01, -4.3494e-01,  8.9682e-02],\n",
      "          [-5.4824e-01,  2.2586e-01,  1.6950e+00]],\n",
      "\n",
      "         [[-9.1739e-01,  1.8871e+00, -1.5179e+00],\n",
      "          [ 1.2003e+00, -5.5797e-01,  2.6736e-01],\n",
      "          [-5.2939e-01, -1.9161e+00,  4.9066e-01]],\n",
      "\n",
      "         [[ 1.3549e+00, -2.7119e-02,  1.7812e+00],\n",
      "          [ 1.0086e+00, -4.2719e-01, -1.7141e+00],\n",
      "          [ 4.6316e-02,  8.7013e-01,  9.7673e-01]],\n",
      "\n",
      "         [[ 4.2659e-01, -1.0866e+00,  1.2207e+00],\n",
      "          [-7.3396e-01,  4.2352e-02, -5.8137e-01],\n",
      "          [ 1.7105e-03, -1.6972e+00,  8.5835e-01]],\n",
      "\n",
      "         [[-8.6449e-03,  8.0620e-01,  4.0657e-01],\n",
      "          [ 7.1347e-01, -4.5871e-01,  9.3667e-01],\n",
      "          [-1.9175e+00, -5.7820e-01,  9.7711e-01]],\n",
      "\n",
      "         [[-1.0672e-01,  1.1908e+00,  8.1571e-02],\n",
      "          [ 2.3109e-01, -2.4445e+00,  1.5317e+00],\n",
      "          [-2.2076e+00, -2.6254e-01, -6.3640e-02]],\n",
      "\n",
      "         [[-2.7079e-01, -9.6244e-02, -7.2355e-01],\n",
      "          [-6.1939e-01, -1.9230e+00, -1.5130e+00],\n",
      "          [ 3.3709e-03,  1.2333e+00, -1.0593e-01]]],\n",
      "\n",
      "\n",
      "        [[[-4.7361e-01,  9.7705e-01,  1.4117e-01],\n",
      "          [-1.8376e+00,  1.3333e+00,  5.4991e-01],\n",
      "          [-7.4934e-01, -6.5864e-01, -7.9952e-01]],\n",
      "\n",
      "         [[-8.5342e-01, -3.8251e-01,  7.0041e-01],\n",
      "          [ 4.6464e-01,  1.7730e+00, -4.9317e-01],\n",
      "          [-2.5031e-01, -9.5274e-01,  4.5490e-01]],\n",
      "\n",
      "         [[-1.5408e-01, -4.0542e-01,  5.5265e-01],\n",
      "          [-2.0821e-01, -9.1971e-01,  2.5784e-01],\n",
      "          [ 5.4357e-01, -1.2070e-01, -1.6519e+00]],\n",
      "\n",
      "         [[ 6.5885e-01, -1.1447e-01, -9.0255e-01],\n",
      "          [ 4.2705e-01, -1.7188e+00,  5.7168e-01],\n",
      "          [ 1.2346e+00,  5.0726e-01, -3.7538e-01]],\n",
      "\n",
      "         [[-4.0448e-01, -5.3865e-01, -3.2126e-01],\n",
      "          [-1.1775e+00,  1.0031e+00,  3.6200e-01],\n",
      "          [-5.1265e-02,  5.9901e-01, -6.8994e-01]],\n",
      "\n",
      "         [[ 3.7802e-01, -1.3381e-01,  6.0202e-01],\n",
      "          [ 3.0395e+00, -3.7553e-01, -9.9873e-02],\n",
      "          [-1.1349e+00,  2.3845e+00,  2.7965e-01]],\n",
      "\n",
      "         [[-9.0218e-01, -1.7182e+00,  2.3859e-03],\n",
      "          [-1.7208e+00, -6.6338e-02, -1.0091e+00],\n",
      "          [ 1.3754e+00, -4.8826e-01,  1.5287e+00]],\n",
      "\n",
      "         [[ 9.3003e-01,  1.3069e+00,  6.6640e-02],\n",
      "          [ 1.7895e+00,  1.1116e-01, -1.2937e+00],\n",
      "          [ 5.9808e-02,  3.8631e-01,  9.9011e-01]],\n",
      "\n",
      "         [[-7.4103e-01, -4.2707e-01,  9.0502e-01],\n",
      "          [-9.5832e-01,  7.8957e-01,  1.0820e+00],\n",
      "          [-2.0352e-01, -3.8323e-01,  1.2363e+00]],\n",
      "\n",
      "         [[ 6.8302e-02,  9.7112e-01,  9.0329e-01],\n",
      "          [-1.0168e+00, -4.9092e-01, -2.2752e-01],\n",
      "          [ 1.5135e-01,  1.1314e-01,  9.4331e-01]]],\n",
      "\n",
      "\n",
      "        [[[ 5.2082e-01,  1.1090e+00,  1.5058e+00],\n",
      "          [-8.5802e-02,  3.0381e-01,  1.7898e+00],\n",
      "          [ 1.7136e-01,  7.1592e-01,  7.0220e-01]],\n",
      "\n",
      "         [[ 2.6729e-02,  1.0858e+00,  7.4080e-01],\n",
      "          [-1.4059e+00, -7.2190e-01,  1.0053e+00],\n",
      "          [-4.4386e-01,  2.2371e-02, -8.2503e-01]],\n",
      "\n",
      "         [[-4.0472e-02, -1.5284e-02,  8.8726e-01],\n",
      "          [-4.7286e-01,  8.7525e-02, -1.3489e+00],\n",
      "          [-4.5784e-02,  5.0426e-01,  1.0426e-01]],\n",
      "\n",
      "         [[ 5.6651e-01, -9.3963e-01, -1.9687e-01],\n",
      "          [ 1.8558e+00,  9.3801e-04,  1.0492e+00],\n",
      "          [ 7.3332e-01, -1.2585e+00,  1.0399e+00]],\n",
      "\n",
      "         [[ 4.8041e-01,  9.1158e-01,  3.7594e-01],\n",
      "          [-1.1657e+00,  1.0596e+00,  8.8196e-01],\n",
      "          [-8.0057e-01, -9.3391e-01,  3.6364e-01]],\n",
      "\n",
      "         [[ 1.6512e+00,  2.2044e+00, -2.7177e-01],\n",
      "          [-2.0609e+00, -1.9627e+00, -3.6775e-01],\n",
      "          [-7.1006e-01,  1.1998e+00, -1.0603e+00]],\n",
      "\n",
      "         [[ 5.4761e-01, -9.5760e-01,  1.1676e+00],\n",
      "          [-1.5552e+00,  1.4128e-01,  1.8889e+00],\n",
      "          [-2.1557e-01, -5.0708e-02,  3.8902e-01]],\n",
      "\n",
      "         [[-8.9169e-01,  2.1059e+00,  1.9103e+00],\n",
      "          [ 5.7192e-01,  1.5531e+00,  1.3912e+00],\n",
      "          [ 1.5677e+00, -9.4109e-01, -1.0711e+00]],\n",
      "\n",
      "         [[ 1.9872e-01, -9.1909e-01,  2.3963e+00],\n",
      "          [-6.8448e-01,  1.0858e+00,  1.8861e+00],\n",
      "          [-3.3736e-01, -5.7978e-01, -1.3313e+00]],\n",
      "\n",
      "         [[ 2.4660e-03, -1.4610e-03,  8.1263e-02],\n",
      "          [ 5.1229e-01, -3.9122e-01,  4.9638e-01],\n",
      "          [-1.7719e-01, -2.5909e-01, -1.5917e+00]]],\n",
      "\n",
      "\n",
      "        [[[-1.0785e+00,  4.1208e-01, -1.6946e+00],\n",
      "          [-1.0967e+00, -3.6528e-01, -7.3403e-01],\n",
      "          [ 1.4480e+00,  9.7005e-01,  9.4975e-01]],\n",
      "\n",
      "         [[ 7.7320e-01,  1.0448e+00, -4.9837e-01],\n",
      "          [ 5.3224e-01,  7.5204e-01, -1.8883e+00],\n",
      "          [ 3.0351e-01,  5.3614e-01, -4.1267e-01]],\n",
      "\n",
      "         [[ 9.1818e-01,  7.3024e-01, -1.1691e+00],\n",
      "          [-6.0581e-01, -6.1149e-02, -1.7445e+00],\n",
      "          [ 1.0414e+00, -8.6846e-02,  5.4014e-01]],\n",
      "\n",
      "         [[-4.4832e-01, -1.4003e+00, -4.2642e-01],\n",
      "          [ 2.5959e-01, -1.0152e-01, -4.7630e-01],\n",
      "          [-1.0077e+00,  5.4072e-02, -4.8994e-01]],\n",
      "\n",
      "         [[-1.4811e+00, -6.0245e-02,  1.6446e-01],\n",
      "          [ 4.7184e-01,  3.0494e-01, -4.2037e-01],\n",
      "          [ 1.7680e+00,  6.7909e-01, -5.3549e-01]],\n",
      "\n",
      "         [[-1.5054e-01,  6.5837e-01,  9.0485e-01],\n",
      "          [ 6.9278e-01, -1.5007e+00,  2.8798e-01],\n",
      "          [ 1.1369e+00, -2.5559e-01,  7.1980e-01]],\n",
      "\n",
      "         [[ 5.9303e-01, -1.1536e+00, -3.8110e-01],\n",
      "          [ 5.5042e-01,  5.3150e-01, -1.7452e+00],\n",
      "          [-5.4097e-01, -2.3603e-02,  3.4605e-02]],\n",
      "\n",
      "         [[-1.0243e+00,  7.5285e-01,  8.1977e-01],\n",
      "          [ 1.7353e+00,  1.8216e+00, -8.5382e-01],\n",
      "          [-2.8199e-01,  4.7911e-01,  2.1427e+00]],\n",
      "\n",
      "         [[ 8.2183e-01, -5.7504e-01,  8.4903e-01],\n",
      "          [ 1.1901e-01,  4.0586e-01,  4.6594e-02],\n",
      "          [-1.7659e+00, -8.5247e-01,  3.0320e-02]],\n",
      "\n",
      "         [[ 6.3853e-01,  2.0515e+00, -5.9265e-02],\n",
      "          [-2.9322e-01, -2.0128e+00,  1.1056e+00],\n",
      "          [-5.0540e-01, -1.8448e-01, -7.8220e-01]]]])\n",
      "conv.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
      "bn.weight tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])\n",
      "bn.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
      "bn.running_mean tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
      "bn.running_var tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])\n",
      "bn.num_batches_tracked tensor(0)\n"
     ]
    }
   ],
   "source": [
    "# Save the model's state dict (parameters only)\n",
    "torch.save(Net.state_dict(), \"net_state_dict.pth\")\n",
    "\n",
    "# Load the state dict\n",
    "Net.load_state_dict(torch.load(\"net_state_dict.pth\"))\n",
    "\n",
    "# Print the loaded state dict\n",
    "for key in Net.state_dict():\n",
    "    print(key, Net.state_dict()[key])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a5b3672",
   "metadata": {},
   "source": [
    "Here, only the model's parameters (state_dict) are saved to disk. When loading, ensure that the network architecture matches the one used during saving to avoid errors.\n",
    "\n",
    "It's important to note that saving the entire model saves the architecture information as well, which can be convenient but may result in larger file sizes. On the other hand, saving only the state_dict is more space-efficient but requires you to recreate the model architecture before loading the parameters.\n",
    "\n",
    "Additionally, you can save and load the state_dict of an optimizer object in a similar manner to maintain the state of optimization algorithms across different sessions."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}