code/utils/optimizer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📝 **Author:** Amirhossein Heydari - 📧 **Email:** <amirhosseinheydari78@gmail.com> - 📍 **Origin:** [mr-pylin/pytorch-workshop](https://github.com/mr-pylin/pytorch-workshop)\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Table of contents**<a id='toc0_'></a>    \n",
    "- [Dependencies](#toc1_)    \n",
    "- [Optimizer](#toc2_)    \n",
    "  - [Built-in Optimizers](#toc2_1_)    \n",
    "    - [Optimizer: SGD](#toc2_1_1_)    \n",
    "    - [Optimizer: Adam](#toc2_1_2_)    \n",
    "    - [Optimizer: Adagrad](#toc2_1_3_)    \n",
    "  - [Custom Optimizers](#toc2_2_)    \n",
    "- [Adjust Learning Rate](#toc3_)    \n",
    "\n",
    "<!-- vscode-jupyter-toc-config\n",
    "\tnumbering=false\n",
    "\tanchor=true\n",
    "\tflat=false\n",
    "\tminLevel=1\n",
    "\tmaxLevel=6\n",
    "\t/vscode-jupyter-toc-config -->\n",
    "<!-- THIS CELL WILL BE REPLACED ON TOC UPDATE. DO NOT WRITE YOUR TEXT IN THIS CELL -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc1_'></a>[Dependencies](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import optim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc2_'></a>[Optimizer](#toc0_)\n",
    "\n",
    "- An **optimizer** updates model **parameters** during training to minimize the loss function.  \n",
    "- It adjusts weights based on **gradients** computed through **backpropagation**.  \n",
    "- Passing `model.parameters()` to the optimizer is **mandatory** to specify which parameters should be updated.  \n",
    "\n",
    "🛠 **Optimizer Methods**:\n",
    "\n",
    "- `zero_grad()` **clears** the **accumulated gradients** from the previous step (**iteration**) to prevent incorrect updates.\n",
    "- `step()` **updates** the model parameters using the computed **gradients**.\n",
    "\n",
    "🖥 **Typical Workflow**:\n",
    "\n",
    "```python\n",
    "output = model(data)            # Forward pass\n",
    "loss = loss_fn(output, target)  # Compute loss\n",
    "loss.backward()                 # Backpropagation (compute gradients)\n",
    "optimizer.step()                # Update parameters\n",
    "optimizer.zero_grad()           # Clear previous gradients\n",
    "```\n",
    "\n",
    "ℹ️ **Learn more**:\n",
    "\n",
    "- details about gradient: [**02-gradient.ipynb**](../02-gradient.ipynb)\n",
    "- details about models: [**model.ipynb**](./model.ipynb)\n",
    "- details about losses: [**loss.ipynb**](./loss.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_1_'></a>[Built-in Optimizers](#toc0_)\n",
    "\n",
    "- PyTorch provides several built-in optimizers in `torch.optim`, each with **different** strategies for updating model parameters.\n",
    "\n",
    "<table style=\"margin:0 auto;\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Optimizer</th>\n",
    "      <th>Best For</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td><pre>optim.SGD()</pre></td>\n",
    "      <td>General training, often best with momentum</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><pre>optim.Adam()</pre></td>\n",
    "      <td>Default choice, works well for most models</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><pre>optim.RMSprop()</pre></td>\n",
    "      <td>RNNs and nonstationary data</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><pre>optim.Adagrad()</pre></td>\n",
    "      <td>Sparse datasets with rare features</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><pre>optim.AdamW()</pre></td>\n",
    "      <td>When weight decay regularization is needed</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- `torch.optim`: [pytorch.org/docs/stable/optim.html](https://pytorch.org/docs/stable/optim.html)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc2_1_1_'></a>[Optimizer: SGD](#toc0_)\n",
    "\n",
    "- Implements stochastic gradient descent (optionally with momentum).\n",
    "- Updates parameters using gradients computed via **backpropagation**. \n",
    "- Works well for convex problems but may struggle with noisy gradients in deep networks.  \n",
    "\n",
    "🎛 **Optimizer Parameters**:\n",
    "\n",
    "<table style=\"margin:0 auto;\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Parameter</th>\n",
    "      <th>Type</th>\n",
    "      <th>Default</th>\n",
    "      <th>Description</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">params</span></td>\n",
    "      <td>iterable</td>\n",
    "      <td><b>Required</b></td>\n",
    "      <td>Iterable of parameters or <span style=\"font-family: monospace;\">named_parameters</span> to optimize. Can also be a list of dicts defining parameter groups.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">lr</span></td>\n",
    "      <td>float or Tensor</td>\n",
    "      <td><span style=\"font-family: monospace;\">1e-3</span></td>\n",
    "      <td>Learning rate (step size).</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">momentum</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">0.0</span></td>\n",
    "      <td>Momentum factor for smoothing updates and accelerating convergence.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">dampening</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">0.0</span></td>\n",
    "      <td>Dampens momentum updates. If <span style=\"font-family: monospace;\">&gt;0</span>, reduces the effect of momentum.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">weight_decay</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">0.0</span></td>\n",
    "      <td>L2 regularization (prevents overfitting).</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">nesterov</span></td>\n",
    "      <td>bool</td>\n",
    "      <td><span style=\"font-family: monospace;\">False</span></td>\n",
    "      <td>Enables <b>Nesterov momentum</b> for faster convergence. Requires <span style=\"font-family: monospace;\">momentum &gt; 0</span>.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">maximize</span></td>\n",
    "      <td>bool</td>\n",
    "      <td><span style=\"font-family: monospace;\">False</span></td>\n",
    "      <td>If <span style=\"font-family: monospace;\">True</span>, maximizes the objective function instead of minimizing it.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">foreach</span></td>\n",
    "      <td>bool or None</td>\n",
    "      <td><span style=\"font-family: monospace;\">None</span></td>\n",
    "      <td>Uses <span style=\"font-family: monospace;\">foreach</span> implementation on CUDA for better performance. Uses more memory but speeds up execution.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">differentiable</span></td>\n",
    "      <td>bool</td>\n",
    "      <td><span style=\"font-family: monospace;\">False</span></td>\n",
    "      <td>If <span style=\"font-family: monospace;\">True</span>, allows autograd to track the optimizer step. Can slow down training.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">fused</span></td>\n",
    "      <td>bool or None</td>\n",
    "      <td><span style=\"font-family: monospace;\">None</span></td>\n",
    "      <td>Uses a fused implementation for <b>better performance</b> on <span style=\"font-family: monospace;\">float32</span>, <span style=\"font-family: monospace;\">float64</span>, <span style=\"font-family: monospace;\">float16</span>, and <span style=\"font-family: monospace;\">bfloat16</span>.</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simple model\n",
    "model = torch.nn.Linear(2, 1)\n",
    "\n",
    "# optimizer\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.01, nesterov=True)\n",
    "\n",
    "# training loop simulation\n",
    "for i in range(3):\n",
    "    loss = model(torch.randn(1, 2)).sum()  # dummy loss\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # log parameters\n",
    "    print(f\"iteration {i+1}/{3}:\")\n",
    "    for n, v in model.named_parameters():\n",
    "        print(f\"{n:6}: {v.data}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc2_1_2_'></a>[Optimizer: Adam](#toc0_)\n",
    "\n",
    "- **Adam (Adaptive Moment Estimation)** Combines **momentum** and **adaptive learning rates** for efficient optimization.  \n",
    "- Uses **exponentially moving averages** of past gradients and squared gradients.  \n",
    "- Works well for non-stationary objectives and sparse gradients.  \n",
    "\n",
    "🎛 **Optimizer Parameters**:  \n",
    "\n",
    "<table style=\"margin:0 auto;\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Parameter</th>\n",
    "      <th>Type</th>\n",
    "      <th>Default</th>\n",
    "      <th>Description</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">params</span></td>\n",
    "      <td>iterable</td>\n",
    "      <td><b>Required</b></td>\n",
    "      <td>Iterable of parameters or <span style=\"font-family: monospace;\">named_parameters</span> to optimize. Can also be a list of dicts defining parameter groups.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">lr</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">1e-3</span></td>\n",
    "      <td>Learning rate (step size).</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">betas</span></td>\n",
    "      <td>Tuple[float, float]</td>\n",
    "      <td><span style=\"font-family: monospace;\">(0.9, 0.999)</span></td>\n",
    "      <td>Coefficients for computing running averages of gradient and squared gradient.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">eps</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">1e-8</span></td>\n",
    "      <td>Term added to the denominator for numerical stability.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">weight_decay</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">0.0</span></td>\n",
    "      <td>L2 regularization (prevents overfitting).</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">amsgrad</span></td>\n",
    "      <td>bool</td>\n",
    "      <td><span style=\"font-family: monospace;\">False</span></td>\n",
    "      <td>Enables AMSGrad variant for better convergence in some cases.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">foreach</span></td>\n",
    "      <td>bool or None</td>\n",
    "      <td><span style=\"font-family: monospace;\">None</span></td>\n",
    "      <td>Uses <span style=\"font-family: monospace;\">foreach</span> implementation on CUDA for better performance. Uses more memory but speeds up execution.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">maximize</span></td>\n",
    "      <td>bool</td>\n",
    "      <td><span style=\"font-family: monospace;\">False</span></td>\n",
    "      <td>If <span style=\"font-family: monospace;\">True</span>, maximizes the objective function instead of minimizing it.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">capturable</span></td>\n",
    "      <td>bool</td>\n",
    "      <td><span style=\"font-family: monospace;\">False</span></td>\n",
    "      <td>Whether this instance is safe to capture in a CUDA graph. Passing <span style=\"font-family: monospace;\">True</span> can impair ungraphed performance, so if you don’t intend to graph capture this instance, leave it <span style=\"font-family: monospace;\">False</span>.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">differentiable</span></td>\n",
    "      <td>bool</td>\n",
    "      <td><span style=\"font-family: monospace;\">False</span></td>\n",
    "      <td>If <span style=\"font-family: monospace;\">True</span>, allows autograd to track the optimizer step. Can slow down training.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">fused</span></td>\n",
    "      <td>bool or None</td>\n",
    "      <td><span style=\"font-family: monospace;\">None</span></td>\n",
    "      <td>Uses a fused implementation for <b>better performance</b> on <span style=\"font-family: monospace;\">float32</span>, <span style=\"font-family: monospace;\">float64</span>, <span style=\"font-family: monospace;\">float16</span>, and <span style=\"font-family: monospace;\">bfloat16</span>.</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simple model\n",
    "model = torch.nn.Linear(2, 1)\n",
    "\n",
    "# optimizer\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.1, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, amsgrad=True)\n",
    "\n",
    "# training loop simulation\n",
    "for i in range(3):\n",
    "    loss = model(torch.randn(1, 2)).sum()  # dummy loss\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # log parameters\n",
    "    print(f\"iteration {i+1}/{3}:\")\n",
    "    for n, v in model.named_parameters():\n",
    "        print(f\"{n:6}: {v.data}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc2_1_3_'></a>[Optimizer: Adagrad](#toc0_)\n",
    "\n",
    "- **Adagrad (Adaptive Gradient Algorithm)** adapts the learning rate individually for each parameter based on past gradients.\n",
    "- Suitable for sparse data and NLP tasks but may suffer from aggressive learning rate decay.\n",
    "\n",
    "🎛 **Optimizer Parameters**:\n",
    "\n",
    "<table>\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Parameter</th>\n",
    "      <th>Type</th>\n",
    "      <th>Default</th>\n",
    "      <th>Description</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">params</span></td>\n",
    "      <td>iterable</td>\n",
    "      <td><b>Required</b></td>\n",
    "      <td>Iterable of parameters or <span style=\"font-family: monospace;\">named_parameters</span> to optimize. Can also be a list of dicts defining parameter groups.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">lr</span></td>\n",
    "      <td>float or Tensor</td>\n",
    "      <td><span style=\"font-family: monospace;\">1e-2</span></td>\n",
    "      <td>Learning rate (step size).</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">lr_decay</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">0.0</span></td>\n",
    "      <td>Decay factor applied to the learning rate.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">weight_decay</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">0.0</span></td>\n",
    "      <td>L2 regularization (prevents overfitting).</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">initial_accumulator_value</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">0.0</span></td>\n",
    "      <td>Initial value of the sum of squares of gradients.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">eps</span></td>\n",
    "      <td>float</td>\n",
    "      <td><span style=\"font-family: monospace;\">1e-10</span></td>\n",
    "      <td>Term added to the denominator to improve numerical stability.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">foreach</span></td>\n",
    "      <td>bool or None</td>\n",
    "      <td><span style=\"font-family: monospace;\">None</span></td>\n",
    "      <td>Uses <span style=\"font-family: monospace;\">foreach</span> implementation on CUDA for better performance. Uses more memory but speeds up execution.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">maximize</span></td>\n",
    "      <td>bool</td>\n",
    "      <td><span style=\"font-family: monospace;\">False</span></td>\n",
    "      <td>If <span style=\"font-family: monospace;\">True</span>, maximizes the objective function instead of minimizing it.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">differentiable</span></td>\n",
    "      <td>bool</td>\n",
    "      <td><span style=\"font-family: monospace;\">False</span></td>\n",
    "      <td>If <span style=\"font-family: monospace;\">True</span>, allows autograd to track the optimizer step. Can slow down training.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">fused</span></td>\n",
    "      <td>bool or None</td>\n",
    "      <td><span style=\"font-family: monospace;\">None</span></td>\n",
    "      <td>Uses a fused implementation for <b>better performance</b> on <span style=\"font-family: monospace;\">float32</span>, <span style=\"font-family: monospace;\">float64</span>, <span style=\"font-family: monospace;\">float16</span>, and <span style=\"font-family: monospace;\">bfloat16</span>. Not supported for sparse or complex gradients.</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simple model\n",
    "model = torch.nn.Linear(2, 1)\n",
    "\n",
    "# optimizer\n",
    "optimizer = optim.Adagrad(\n",
    "    model.parameters(), lr=0.1, lr_decay=0.01, weight_decay=0.01, initial_accumulator_value=0.1, eps=1e-10\n",
    ")\n",
    "\n",
    "# training loop simulation\n",
    "for i in range(3):\n",
    "    loss = model(torch.randn(1, 2)).sum()  # dummy loss\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # log parameters\n",
    "    print(f\"iteration {i+1}/{3}:\")\n",
    "    for n, v in model.named_parameters():\n",
    "        print(f\"{n:6}: {v.data}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_2_'></a>[Custom Optimizers](#toc0_)\n",
    "\n",
    "- PyTorch allows defining **custom optimizers** by extending `torch.optim.Optimizer`.  \n",
    "- Custom optimizers give full control over parameter updates, allowing modifications beyond standard methods like **SGD**, **Adam**, or **Adagrad**.\n",
    "- To implement a custom optimizer:\n",
    "  1. **Inherit** from `torch.optim.Optimizer`.\n",
    "  2. **Initialize** parameters and defaults in `__init__()`.\n",
    "  3. **Implement `step()`**, which defines how parameters are updated each iteration.\n",
    "\n",
    "**Defining `step()`**\n",
    "- The `step()` function is where the **gradient-based update rule** is applied.\n",
    "- Inside `step()`, iterate over `self.param_groups` and update each parameter using its gradient.\n",
    "- Use `@torch.no_grad()` to disable gradient tracking during updates.\n",
    "- Support an optional `closure` function for loss recomputation (useful in some optimizers like **LBFGS**).\n",
    "\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- `torch.optim.Optimizer`: [pytorch.org/docs/stable/optim.html#torch.optim.Optimizer](https://pytorch.org/docs/stable/optim.html)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomSGD(optim.Optimizer):\n",
    "    def __init__(self, params, lr=0.01):\n",
    "        if lr < 0.0:\n",
    "            raise ValueError(f\"Invalid learning rate: {lr}\")\n",
    "        defaults = {\"lr\": lr}\n",
    "        super().__init__(params, defaults)\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def step(self, closure=None):\n",
    "        loss = None\n",
    "        if closure is not None:\n",
    "            loss = closure()\n",
    "\n",
    "        for group in self.param_groups:\n",
    "            lr = group[\"lr\"]\n",
    "            for param in group[\"params\"]:\n",
    "                if param.grad is not None:\n",
    "                    param -= lr * param.grad\n",
    "\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simple model\n",
    "model = torch.nn.Linear(2, 1)\n",
    "\n",
    "# optimizer\n",
    "optimizer = CustomSGD(model.parameters(), lr=0.1)\n",
    "\n",
    "# training loop simulation\n",
    "for i in range(3):\n",
    "    loss = model(torch.randn(1, 2)).sum()  # dummy loss\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # log parameters\n",
    "    print(f\"iteration {i+1}/{3}:\")\n",
    "    for n, v in model.named_parameters():\n",
    "        print(f\"{n:6}: {v.data}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc3_'></a>[Adjust Learning Rate](#toc0_)\n",
    "\n",
    "- PyTorch provides **learning rate schedulers** in `torch.optim.lr_scheduler` to dynamically adjust the learning rate during training.\n",
    "- Learning rate scheduling should be applied **after** optimizer’s update.\n",
    "- Learning rate schedulers **must be called per epoch or per iteration**, depending on the chosen strategy.\n",
    "- Common Learning Rate Schedulers:\n",
    "<table style=\"margin:0 auto;\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Scheduler</th>\n",
    "      <th>Description</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">StepLR</span></td>\n",
    "      <td>Decays LR every <span style=\"font-family: monospace;\">step_size</span> epochs.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">MultiStepLR</span></td>\n",
    "      <td>Decays LR at multiple predefined epochs.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">ExponentialLR</span></td>\n",
    "      <td>Decays LR by a fixed multiplicative factor every epoch.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">ReduceLROnPlateau</span></td>\n",
    "      <td>Reduces LR when a metric has stopped improving.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">CosineAnnealingLR</span></td>\n",
    "      <td>Uses a cosine function for annealing LR.</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td><span style=\"font-family: monospace;\">LambdaLR</span></td>\n",
    "      <td>Custom LR scheduling via a user-defined function.</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- How to adjust learning rate: [pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate](https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, \"min\", patience=3)"
   ]
  }
 ],
 "metadata": {
  "author_email": "AmirhosseinHeydari78@gmail.com",
  "author_github": "https://github.com/mr-pylin",
  "author_name": "Amirhossein Heydari",
  "kernelspec": {
   "display_name": "pytorch-workshop-U_zYfVTd-py3.12",
   "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.12.8"
  },
  "origin_repo": "https://github.com/mr-pylin/pytorch-workshop"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}