03_Encapsulate_your_model_with_Modules.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5e172911",
   "metadata": {},
   "source": [
    "For enhancing code organization, it's advisable to employ PyTorch's modules, which serve as structured containers for parameters while encapsulating model operations. Modules offer a systematic approach to managing various aspects of a neural network, providing clarity and ease of maintenance. As an illustration, consider a scenario where you aim to depict a linear model represented by the equation y = ax + b. By leveraging PyTorch's modules, you can implement this model with precision and readability, as exemplified below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d548f9a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "class Net(torch.nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.a = torch.nn.Parameter(torch.rand(1))\n",
    "        self.b = torch.nn.Parameter(torch.rand(1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        yhat = self.a * x + self.b\n",
    "        return yhat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96c400d9-ca95-42f9-bb74-4f8b6fc4fa6a",
   "metadata": {},
   "source": [
    "To use this model in practice you instantiate the module and simply call it like a function:\r\n",
    "\r\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "233896a7-090b-4c10-b0db-d6d123eeaa29",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.arange(100, dtype=torch.float32)\n",
    "\n",
    "net = Net()\n",
    "y = net(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9d46b65-2a30-4a82-9e83-0771616a3ed2",
   "metadata": {},
   "source": [
    "Parameters are essentially tensors with `requires_grad` set to true. It's convenient to use parameters because you can simply retrieve them all with module's `parameters()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "83ee32b5-48ae-45fb-ac8e-30d30326678f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([0.9216], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0.6813], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "for p in net.parameters():\n",
    "    print(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30ff34d4-0b82-4bcc-8377-1ae7ca7dcaad",
   "metadata": {},
   "source": [
    "Now, say you have an unknown function `y = 5x + 3 + some noise`, and you want to optimize the parameters of your model to fit this function. You can start by sampling some points from your function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0005803d-c61e-4d64-9d93-f03242da2d90",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.arange(100, dtype=torch.float32) / 100\n",
    "y = 5 * x + 3 + torch.rand(100) * 0.3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0e48c2e-977f-4f7d-855b-36fe7a93539b",
   "metadata": {},
   "source": [
    "Similar to the previous example, you can define a loss function and optimize the parameters of your model as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dc212ce0-027c-4bc2-be19-d963d8440103",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([4.9650], requires_grad=True) Parameter containing:\n",
      "tensor([3.1714], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.SGD(net.parameters(), lr=0.01)\n",
    "\n",
    "for i in range(10000):\n",
    "  net.zero_grad()\n",
    "  yhat = net(x)\n",
    "  loss = criterion(yhat, y)\n",
    "  loss.backward()\n",
    "  optimizer.step()\n",
    "\n",
    "print(net.a, net.b) # Should be close to 5 and 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f00de16-e664-49f9-af5a-e575e2595c0c",
   "metadata": {},
   "source": [
    "PyTorch comes with a number of predefined modules. One such module is `torch.nn.Linear` which is a more general form of a linear function than what we defined above. We can rewrite our module above using `torch.nn.Linear` like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "29266c6b-ab9a-43e9-8f57-73425702c4d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(torch.nn.Module):\n",
    "\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.linear = torch.nn.Linear(1, 1)\n",
    "\n",
    "  def forward(self, x):\n",
    "    yhat = self.linear(x.unsqueeze(1)).squeeze(1)\n",
    "    return yhat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9950578d-29c4-4efd-a794-5d97782bfdf7",
   "metadata": {},
   "source": [
    "Note that we used squeeze and unsqueeze since `torch.nn.Linear` operates on batch of vectors as opposed to scalars.\n",
    "\n",
    "By default calling parameters() on a module will return the parameters of all its submodules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c3539bd9-6d09-49a1-941a-1ad63c4f92ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([[-0.7680]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0.4553], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "net = Net()\n",
    "for p in net.parameters():\n",
    "    print(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "185c534f-81de-4dd6-8b02-8a26b8372343",
   "metadata": {},
   "source": [
    "There are some predefined modules that act as a container for other modules. The most commonly used container module is `torch.nn.Sequential`. As its name implies it's used to to stack multiple modules (or layers) on top of each other. For example to stack two Linear layers with a `ReLU` nonlinearity in between you can do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4019cfa4-2757-479a-8300-b856db4cbc5d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}