05_Gradient_Optimization/16_Smoothing_Gradients_With_Gradient_Penalty.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9e0c293c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "903f8943",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import dependencies\n",
    "import math\n",
    "from collections import OrderedDict\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "from torchvision.datasets.mnist import MNIST\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.sampler import RandomSampler, SequentialSampler\n",
    "import torch.nn.init as init\n",
    "\n",
    "import random\n",
    "import numpy as np\n",
    "import psutil\n",
    "from torchsummary import summary\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import warnings \n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4b56f9ea-54b0-4175-972a-523290fdbe41",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random seed set to 42\n"
     ]
    }
   ],
   "source": [
    "def set_experiment_seed(seed=42):\n",
    "    \"\"\"\n",
    "    Set random seeds and CUDA-related flags for experiment reproducibility.\n",
    "    \"\"\"\n",
    "    # Set random seeds for reproducibility\n",
    "    torch.manual_seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    random.seed(seed)\n",
    "\n",
    "    # If using GPU, set random seed for CUDA operations\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.manual_seed(seed)\n",
    "        torch.backends.cudnn.deterministic = True\n",
    "        torch.backends.cudnn.benchmark = False  # Set to False for reproducibility\n",
    "\n",
    "    # Optionally log seed information\n",
    "    print(f\"Random seed set to {seed}\")\n",
    "\n",
    "set_experiment_seed()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "079ce688",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download MNIST dataset\n",
    "data_train = MNIST(\n",
    "    \"./data/mnist\", download=True, \n",
    "    transform=transforms.Compose([\n",
    "        transforms.Resize((32, 32)),\n",
    "        transforms.RandomRotation(degrees=15),\n",
    "        transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),\n",
    "        transforms.RandomVerticalFlip(),\n",
    "        transforms.ColorJitter(brightness=0.2, contrast=0.2),\n",
    "        transforms.GaussianBlur(kernel_size=1),\n",
    "        transforms.ToTensor()\n",
    "    ])\n",
    ")\n",
    "        \n",
    "data_test = MNIST(\n",
    "    \"./data/mnist\", download=True, train=False,\n",
    "    transform=transforms.Compose([\n",
    "        transforms.Resize((32, 32)),\n",
    "        transforms.ToTensor()\n",
    "    ])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "372c5b4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define hyperparameters\n",
    "N_CLASSES = 10\n",
    "EPOCHS = 20\n",
    "BATCH_SIZE = 512\n",
    "DROPOUT_PROB = 0.2\n",
    "USE_BATCHNORM = True\n",
    "LABEL_SMOOTHING = 0.1\n",
    "LR = 1.34E-03\n",
    "WARMUP_PROPORTION = 0.1\n",
    "WEIGHT_DECAY = 1e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "93918f18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.figure(figsize=(10,10))\n",
    "# for i in range(25):\n",
    "#     plt.subplot(5,5,i+1)\n",
    "#     plt.xticks([])\n",
    "#     plt.yticks([])\n",
    "#     plt.grid(False)\n",
    "#     plt.imshow(data_train[i][0].permute(1,2,0).numpy(), cmap=plt.cm.binary)\n",
    "#     plt.xlabel(data_train[i][1])\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce264b0f-c32e-4907-bc34-c7644df7d2c7",
   "metadata": {},
   "source": [
    "`torch.utils.data.DataLoader` supports asynchronous data loading and data augmentation in separate worker subprocesses. The default setting for DataLoader is `num_workers=0`, which means that the data loading is synchronous and done in the main process. As a result the main training process has to wait for the data to be available to continue the execution.\r\n",
    "\r\n",
    "Settin`g num_workers >` 0 enables asynchronous data loading and overlap between the training and data loading. num_workers should be tuned depending on the workload, CPU, GPU, and location of training data.`\r\n",
    "\r\n",
    "DataLo`ader acce`pts pin_me`mory argument, which defaults` to F`alse. When using a GPU it’s better to `set pin_memory=`True, this instru`cts DataLo`ader to use pinned memory and enables faster and asynchronous memory copy from the host to the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fb249649-163b-495a-bd21-94711d5ec411",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_num_workers():\n",
    "    # Calculate the number of CPU cores\n",
    "    num_cpu_cores = os.cpu_count()\n",
    "\n",
    "    # Set a safe maximum multiplier value (e.g., 0.5) to avoid using all available resources\n",
    "    max_multiplier = 0.5\n",
    "\n",
    "    # Calculate the number of workers based on the available resources\n",
    "    multiplier = min(sum(psutil.cpu_percent(interval=1, percpu=True)) / 100.0, max_multiplier)\n",
    "    \n",
    "    # Ensure that num_workers is at least 1\n",
    "    num_workers = max(1, int(num_cpu_cores * multiplier))\n",
    "\n",
    "    return num_workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5b752015",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define train dataloader\n",
    "train_dataloader = DataLoader(\n",
    "    data_train, batch_size=BATCH_SIZE, \n",
    "    sampler=RandomSampler(data_train), \n",
    "    pin_memory=True, num_workers=calculate_num_workers()\n",
    ")\n",
    "\n",
    "# Define test dataloader\n",
    "test_dataloader = DataLoader(\n",
    "    data_test, batch_size=BATCH_SIZE, \n",
    "    sampler=SequentialSampler(data_test),\n",
    "    pin_memory=True, num_workers=calculate_num_workers()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6864f541-00b1-4851-a1da-288992d8d709",
   "metadata": {},
   "source": [
    "`torch.nn.Conv2d()` has `bias` parameter which defaults to `True` (the same is true for `Conv1d` and `Conv3d` ).\n",
    "\n",
    "If a `nn.Conv2d` layer is directly followed by a `nn.BatchNorm2d` layer, then the bias in the convolution is not needed, instead use `nn.Conv2d(..., bias=False, ....)`. Bias is not needed because in the first step `BatchNorm` subtracts the mean, which effectively cancels out the effect of bias.\n",
    "\n",
    "This is also applicable to 1d and 3d convolutions as long as `BatchNorm`(or other normalization layer) normalizes on the same dimension as convolution’s bias."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "35af50e0-5861-444c-ab61-8581c61f9af0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LeNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LeNet, self).__init__()\n",
    "                \n",
    "        self.features1 = nn.Sequential(\n",
    "            nn.Conv2d(1, 6, kernel_size=(5, 5), bias=not USE_BATCHNORM),\n",
    "            self._get_norm_layer(6),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=(2, 2), stride=2),\n",
    "            nn.Dropout2d(p=DROPOUT_PROB)\n",
    "        )\n",
    "        \n",
    "        self.features2 = nn.Sequential(\n",
    "            nn.Conv2d(6, 16, kernel_size=(5, 5), bias=not USE_BATCHNORM),\n",
    "            self._get_norm_layer(16),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=(2, 2), stride=2),\n",
    "            nn.Dropout2d(p=DROPOUT_PROB)\n",
    "        )\n",
    "\n",
    "        self.features3 = nn.Sequential(\n",
    "            nn.Conv2d(16, 120, kernel_size=(5, 5), bias=not USE_BATCHNORM),\n",
    "            self._get_norm_layer(120),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout2d(p=DROPOUT_PROB)\n",
    "        )\n",
    "\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(120, 84, bias=True),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=DROPOUT_PROB),\n",
    "            nn.Linear(84, 10, bias=True)\n",
    "        )\n",
    "\n",
    "        # Initialize layers\n",
    "        self.apply(self._initialize_weights)\n",
    "\n",
    "    def _get_norm_layer(self, channels):\n",
    "        if USE_BATCHNORM:\n",
    "            return nn.BatchNorm2d(channels)\n",
    "        else:\n",
    "            return nn.Identity()\n",
    "\n",
    "    def _initialize_weights(self, m):\n",
    "        if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):\n",
    "            init.kaiming_uniform_(m.weight)\n",
    "            if m.bias is not None:\n",
    "                init.constant_(m.bias, 0)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.features1(x)\n",
    "        x = self.features2(x)\n",
    "        x = self.features3(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.classifier(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "621c6db2-e8a3-42cc-ab36-5694abdf5764",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "31f61bca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1           [512, 6, 28, 28]             150\n",
      "       BatchNorm2d-2           [512, 6, 28, 28]              12\n",
      "              ReLU-3           [512, 6, 28, 28]               0\n",
      "         MaxPool2d-4           [512, 6, 14, 14]               0\n",
      "         Dropout2d-5           [512, 6, 14, 14]               0\n",
      "            Conv2d-6          [512, 16, 10, 10]           2,400\n",
      "       BatchNorm2d-7          [512, 16, 10, 10]              32\n",
      "              ReLU-8          [512, 16, 10, 10]               0\n",
      "         MaxPool2d-9            [512, 16, 5, 5]               0\n",
      "        Dropout2d-10            [512, 16, 5, 5]               0\n",
      "           Conv2d-11           [512, 120, 1, 1]          48,000\n",
      "      BatchNorm2d-12           [512, 120, 1, 1]             240\n",
      "             ReLU-13           [512, 120, 1, 1]               0\n",
      "        Dropout2d-14           [512, 120, 1, 1]               0\n",
      "           Linear-15                  [512, 84]          10,164\n",
      "             ReLU-16                  [512, 84]               0\n",
      "          Dropout-17                  [512, 84]               0\n",
      "           Linear-18                  [512, 10]             850\n",
      "================================================================\n",
      "Total params: 61,848\n",
      "Trainable params: 61,848\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 2.00\n",
      "Forward/backward pass size (MB): 89.09\n",
      "Params size (MB): 0.24\n",
      "Estimated Total Size (MB): 91.32\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# Instantiate the model\n",
    "model = LeNet()\n",
    "model = model.to(device)\n",
    "\n",
    "# Test with a random input\n",
    "summary(model, input_size=(1, 32, 32), batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d186351-babd-4d4e-886e-d9338ca71640",
   "metadata": {},
   "source": [
    "#### 1. Cross Entropy Loss\n",
    "The standard cross-entropy loss for classification tasks is given by:\n",
    "\n",
    "$$ \\text{Traditional Cross Entropy Loss: } H(y, \\hat{y}) = - \\sum_i y_i \\log(\\hat{y}_i) $$\n",
    "\n",
    " - $y_i$ is a binary indicator of whether class $i$ is the correct classification.  \n",
    " - $p_i$ is the predicted probability of class $i$.\n",
    "\n",
    "#### 2. Label Smoothed Cross Entropy\n",
    "Label Smoothing Cross Entropy Loss introduces a modification to the target distribution:\n",
    "\n",
    " $$ \\text{Label Smoothed Cross Entropy Loss} = - \\sum_i \\left( (1 - \\text{smoothing}) \\cdot 1_{\\{y_i\\}} + \\frac{\\text{smoothing}}{C-1} \\cdot 1_{\\{1 - y_i\\}} \\right) \\cdot \\log(p_i) $$\n",
    "\n",
    "Where:\n",
    "- $1_{\\{y_i\\}}$ is a binary indicator of whether class $i$ is the correct classification.\n",
    "- $p_i$ is the predicted probability of class $i$.\n",
    "- $C$ is the number of classes.\n",
    "- $\\text{smoothing}$  is the smoothing factor.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5556bcc8-3ff2-4169-b316-0a7b17aa1d2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LabelSmoothedCrossEntropy(nn.Module):\n",
    "    def __init__(self, num_classes, smoothing):\n",
    "        super(LabelSmoothedCrossEntropy, self).__init__()\n",
    "        self.num_classes = num_classes\n",
    "        self.smoothing = smoothing\n",
    "        self.confidence = 1.0 - smoothing\n",
    "\n",
    "    def forward(self, input_logits, target):\n",
    "        log_probs = F.log_softmax(input_logits, dim=-1)\n",
    "        true_dist = torch.zeros_like(log_probs)\n",
    "        true_dist.fill_(self.smoothing / (self.num_classes - 1))\n",
    "        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)\n",
    "        loss = -torch.sum(true_dist * log_probs) / input_logits.size(0)\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "262d50c6-de05-416f-8ad7-14792edf403c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize loss\n",
    "# criterion = nn.CrossEntropyLoss()\n",
    "criterion = LabelSmoothedCrossEntropy(num_classes=N_CLASSES, smoothing=LABEL_SMOOTHING)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1365bb52-8cf9-414e-bb7e-d87c073b933b",
   "metadata": {},
   "source": [
    "#### Adam w/ Weight Decay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3412026a-4b8b-4070-9781-6efb964d3446",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize optimizer\n",
    "# optimizer = optim.Adam(model.parameters(), lr=LR)\n",
    "optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f78d5cf-5ad7-4d10-be31-cb8d498227c4",
   "metadata": {},
   "source": [
    "#### Cosine Annealing With Warmup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6ebc561a-59b6-4848-92f0-bb32d4682d1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lr_lambda(initial_lr, warmup_steps, total_steps):\n",
    "    def lr_lambda(current_step):\n",
    "        if current_step < warmup_steps:\n",
    "            return initial_lr + (1.0 - initial_lr) * float(current_step) / float(max(1, warmup_steps))\n",
    "        else:\n",
    "            return max(0.0, 0.5 * (1.0 + math.cos(math.pi * (current_step - warmup_steps) / float(total_steps - warmup_steps))))\n",
    "    return lr_lambda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b45993ef-a7a2-4154-9eeb-d10e5e003e4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "len_dataloader = len(train_dataloader)\n",
    "num_update_steps_per_epoch = max(len_dataloader, 1) \n",
    "num_examples = len(train_dataloader.dataset)\n",
    "max_steps = math.ceil(EPOCHS * num_update_steps_per_epoch)\n",
    "num_warmup_steps = math.ceil(max_steps * WARMUP_PROPORTION)\n",
    "lr_lambda = get_lr_lambda(LR, num_warmup_steps, max_steps)\n",
    "scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab52c285-645c-4c4f-a051-565c8eb72c94",
   "metadata": {},
   "source": [
    "#### Adaptive Gradient Clipping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d4216258-a9d0-4eda-9aa0-e91543d2f4d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def unitwise_norm(x, norm_type=2.0):\n",
    "    if x.ndim <= 1:\n",
    "        return x.norm(norm_type)\n",
    "    else:\n",
    "        return x.norm(norm_type, dim=tuple(range(1, x.ndim)), keepdim=True)\n",
    "\n",
    "\n",
    "def adaptive_clip_grad(parameters, clip_factor=0.01, eps=1e-3, norm_type=2.0):\n",
    "    if isinstance(parameters, torch.Tensor):\n",
    "        parameters = [parameters]\n",
    "    for p in parameters:\n",
    "        if p.grad is None:\n",
    "            continue\n",
    "        p_data = p.detach()\n",
    "        g_data = p.grad.detach()\n",
    "        max_norm = unitwise_norm(p_data, norm_type=norm_type).clamp_(min=eps).mul_(clip_factor)\n",
    "        grad_norm = unitwise_norm(g_data, norm_type=norm_type)\n",
    "        clipped_grad = g_data * (max_norm / grad_norm.clamp(min=1e-6))\n",
    "        new_grads = torch.where(grad_norm < max_norm, g_data, clipped_grad)\n",
    "        p.grad.detach().copy_(new_grads)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e88f5212-0ec9-4be0-bca4-260d4d6f6cf1",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "#### LR Finder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c872fb49-b472-44e3-bc6c-bfced9e7a162",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainDataLoaderIter:\n",
    "    \"\"\"\n",
    "    Iterator for a PyTorch DataLoader, facilitating easy iteration over batches of data.\n",
    "\n",
    "    Args:\n",
    "        dataloader (torch.utils.data.DataLoader): The PyTorch DataLoader providing batches of data.\n",
    "        auto_reset (bool, optional): If True, the iterator resets to the beginning of the DataLoader\n",
    "            when it reaches the end. If False, a StopIteration exception is raised at the end.\n",
    "            Default is True.\n",
    "\n",
    "    Methods:\n",
    "        inputs_labels_from_batch(batch_data):\n",
    "            Extracts inputs and labels from a batch of data.\n",
    "\n",
    "        __next__():\n",
    "            Gets the next batch of data from the DataLoader, automatically resetting if necessary.\n",
    "\n",
    "        __iter__():\n",
    "            Returns the iterator object.\n",
    "\n",
    "    Usage:\n",
    "        # Create a TrainDataLoaderIter object\n",
    "        train_iter = TrainDataLoaderIter(train_dataloader)\n",
    "\n",
    "        # Iterate over batches of data\n",
    "        for inputs, labels in train_iter:\n",
    "            # Perform training using the current batch\n",
    "\n",
    "    Example with custom inputs_labels_from_batch:\n",
    "        class CustomTrainDataLoaderIter(TrainDataLoaderIter):\n",
    "            def inputs_labels_from_batch(self, batch_data):\n",
    "                # Custom logic to extract inputs and labels from the batch_data\n",
    "                inputs, labels = batch_data['data'], batch_data['label']\n",
    "                return inputs, labels\n",
    "    \"\"\"\n",
    "    def __init__(self, dataloader, auto_reset=True):\n",
    "        \"\"\"\n",
    "        Initialize the TrainDataLoaderIter.\n",
    "\n",
    "        Args:\n",
    "            dataloader (torch.utils.data.DataLoader): The PyTorch DataLoader providing batches of data.\n",
    "            auto_reset (bool, optional): If True, the iterator resets to the beginning of the DataLoader\n",
    "                when it reaches the end. If False, a StopIteration exception is raised at the end.\n",
    "                Default is True.\n",
    "        \"\"\"\n",
    "        self.dataloader = dataloader\n",
    "        self._iterator = iter(dataloader)\n",
    "        self.auto_reset = auto_reset\n",
    "\n",
    "    def inputs_labels_from_batch(self, batch_data):\n",
    "        \"\"\"\n",
    "        Extract inputs and labels from a batch of data.\n",
    "\n",
    "        Args:\n",
    "            batch_data: Batch of data returned by the DataLoader.\n",
    "\n",
    "        Returns:\n",
    "            Tuple: A tuple containing inputs and labels extracted from the batch_data.\n",
    "        \"\"\"\n",
    "        inputs, labels = batch_data\n",
    "        return inputs, labels\n",
    "\n",
    "    def __next__(self):\n",
    "        \"\"\"\n",
    "        Get the next batch of data from the DataLoader, automatically resetting if necessary.\n",
    "\n",
    "        Returns:\n",
    "            Tuple: A tuple containing inputs and labels for the next batch of data.\n",
    "\n",
    "        Raises:\n",
    "            StopIteration: If auto_reset is False and the end of the DataLoader is reached.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            batch = next(self._iterator)\n",
    "            inputs, labels = self.inputs_labels_from_batch(batch)\n",
    "        except StopIteration:\n",
    "            if self.auto_reset:\n",
    "                self._iterator = iter(self.dataloader)\n",
    "                batch = next(self._iterator)\n",
    "                inputs, labels = self.inputs_labels_from_batch(batch)\n",
    "            else:\n",
    "                raise\n",
    "\n",
    "        return inputs, labels\n",
    "\n",
    "    def __iter__(self):\n",
    "        \"\"\"\n",
    "        Returns the iterator object.\n",
    "\n",
    "        Returns:\n",
    "            self: The iterator object.\n",
    "        \"\"\"\n",
    "        return self"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b37eec82-ec8f-407e-be48-d3c532047c7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearLR(torch.optim.lr_scheduler._LRScheduler):\n",
    "    \"\"\"\n",
    "    Linear learning rate scheduler.\n",
    "\n",
    "    Args:\n",
    "        optimizer (torch.optim.Optimizer): Optimizer to adjust the learning rate for.\n",
    "        end_lr (float): The final learning rate after the specified number of iterations.\n",
    "        num_iter (int): The total number of iterations for the learning rate schedule.\n",
    "        last_epoch (int, optional): The index of the last epoch. Default is -1.\n",
    "\n",
    "    Methods:\n",
    "        get_lr():\n",
    "            Calculate the linearly adjusted learning rates based on the current iteration.\n",
    "\n",
    "    Usage:\n",
    "        # Create a LinearLR scheduler\n",
    "        linear_scheduler = LinearLR(optimizer, end_lr=0.001, num_iter=100)\n",
    "\n",
    "        # Inside the training loop\n",
    "        for epoch in range(num_epochs):\n",
    "            for batch in data_loader:\n",
    "                # Perform training using the current batch\n",
    "\n",
    "                # Update the learning rate\n",
    "                linear_scheduler.step()\n",
    "\n",
    "    Note: Ensure that the scheduler is called within the training loop after each training step.\n",
    "\n",
    "    Example:\n",
    "        linear_scheduler = LinearLR(optimizer, end_lr=0.001, num_iter=100)\n",
    "        for epoch in range(num_epochs):\n",
    "            for batch in data_loader:\n",
    "                # Perform training using the current batch\n",
    "\n",
    "                # Update the learning rate\n",
    "                linear_scheduler.step()\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, optimizer, end_lr, num_iter, last_epoch=-1):\n",
    "        \"\"\"\n",
    "        Initialize the LinearLR scheduler.\n",
    "\n",
    "        Args:\n",
    "            optimizer (torch.optim.Optimizer): Optimizer to adjust the learning rate for.\n",
    "            end_lr (float): The final learning rate after the specified number of iterations.\n",
    "            num_iter (int): The total number of iterations for the learning rate schedule.\n",
    "            last_epoch (int, optional): The index of the last epoch. Default is -1.\n",
    "        \"\"\"\n",
    "        self.end_lr = end_lr\n",
    "        self.num_iter = num_iter\n",
    "        super(LinearLR, self).__init__(optimizer, last_epoch)\n",
    "\n",
    "    def get_lr(self):\n",
    "        \"\"\"\n",
    "        Calculate the linearly adjusted learning rates based on the current iteration.\n",
    "\n",
    "        Returns:\n",
    "            list: List of adjusted learning rates for each parameter group.\n",
    "        \"\"\"\n",
    "        r = self.last_epoch / (self.num_iter - 1)\n",
    "        return [base_lr + r * (self.end_lr - base_lr) for base_lr in self.base_lrs]\n",
    "\n",
    "    \n",
    "class ExponentialLR(torch.optim.lr_scheduler._LRScheduler):\n",
    "    \"\"\"\n",
    "    Exponential learning rate scheduler.\n",
    "\n",
    "    Args:\n",
    "        optimizer (torch.optim.Optimizer): Optimizer to adjust the learning rate for.\n",
    "        end_lr (float): The final learning rate after the specified number of iterations.\n",
    "        num_iter (int): The total number of iterations for the learning rate schedule.\n",
    "        last_epoch (int, optional): The index of the last epoch. Default is -1.\n",
    "\n",
    "    Methods:\n",
    "        get_lr():\n",
    "            Calculate the exponentially adjusted learning rates based on the current iteration.\n",
    "\n",
    "    Usage:\n",
    "        # Create an ExponentialLR scheduler\n",
    "        exponential_scheduler = ExponentialLR(optimizer, end_lr=0.001, num_iter=100)\n",
    "\n",
    "        # Inside the training loop\n",
    "        for epoch in range(num_epochs):\n",
    "            for batch in data_loader:\n",
    "                # Perform training using the current batch\n",
    "\n",
    "                # Update the learning rate\n",
    "                exponential_scheduler.step()\n",
    "\n",
    "    Note: Ensure that the scheduler is called within the training loop after each training step.\n",
    "\n",
    "    Example:\n",
    "        exponential_scheduler = ExponentialLR(optimizer, end_lr=0.001, num_iter=100)\n",
    "        for epoch in range(num_epochs):\n",
    "            for batch in data_loader:\n",
    "                # Perform training using the current batch\n",
    "\n",
    "                # Update the learning rate\n",
    "                exponential_scheduler.step()\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, optimizer, end_lr, num_iter, last_epoch=-1):\n",
    "        \"\"\"\n",
    "        Initialize the ExponentialLR scheduler.\n",
    "\n",
    "        Args:\n",
    "            optimizer (torch.optim.Optimizer): Optimizer to adjust the learning rate for.\n",
    "            end_lr (float): The final learning rate after the specified number of iterations.\n",
    "            num_iter (int): The total number of iterations for the learning rate schedule.\n",
    "            last_epoch (int, optional): The index of the last epoch. Default is -1.\n",
    "        \"\"\"\n",
    "        self.end_lr = end_lr\n",
    "        self.num_iter = num_iter\n",
    "        super(ExponentialLR, self).__init__(optimizer, last_epoch)\n",
    "\n",
    "    def get_lr(self):\n",
    "        \"\"\"\n",
    "        Calculate the exponentially adjusted learning rates based on the current iteration.\n",
    "\n",
    "        Returns:\n",
    "            list: List of adjusted learning rates for each parameter group.\n",
    "        \"\"\"\n",
    "        r = self.last_epoch / (self.num_iter - 1)\n",
    "        return [base_lr * (self.end_lr / base_lr) ** r for base_lr in self.base_lrs]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1bde5469-f2cd-456e-b904-2bbff34dbdf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LearningRateFinder:\n",
    "    \"\"\"\n",
    "    Learning rate finder class for determining an optimal learning rate during model training.\n",
    "\n",
    "    Args:\n",
    "        model (torch.nn.Module): The neural network model to train.\n",
    "        criterion: The loss function used for training.\n",
    "        optimizer: The optimizer used for updating model parameters.\n",
    "        device (torch.device): The device (CPU or GPU) on which to perform training.\n",
    "\n",
    "    Attributes:\n",
    "        model (torch.nn.Module): The neural network model to train.\n",
    "        criterion: The loss function used for training.\n",
    "        optimizer: The optimizer used for updating model parameters.\n",
    "        device (torch.device): The device (CPU or GPU) on which to perform training.\n",
    "        history (dict): A dictionary to store the learning rate and loss history during search.\n",
    "            Keys: 'lr' for learning rates, 'loss' for corresponding losses.\n",
    "        best_loss: The best observed loss during the learning rate search.\n",
    "\n",
    "    Methods:\n",
    "        find_lr(train_dataloader, start_lr=None, end_lr=10, num_iter=100, step_mode=\"exp\",\n",
    "                smooth_f=0.05, diverge_th=5, early_stop_patience=None):\n",
    "            Search for the optimal learning rate within the specified range.\n",
    "\n",
    "        _train_batch(train_iter):\n",
    "            Train a single batch and update the model parameters.\n",
    "\n",
    "        set_learning_rate(start_lrs):\n",
    "            Set the learning rate(s) for the optimizer.\n",
    "\n",
    "        plot(skip_start=10, skip_end=5, log_lr=True, show_lr=None):\n",
    "            Plot the learning rate search results and suggest an optimal learning rate.\n",
    "\n",
    "    Usage:\n",
    "        # Create a LearningRateFinder instance\n",
    "        lr_finder = LearningRateFinder(model, criterion, optimizer, device)\n",
    "\n",
    "        # Find the optimal learning rate\n",
    "        lr_finder.find_lr(train_dataloader, start_lr=1e-5, end_lr=10, num_iter=100)\n",
    "\n",
    "        # Plot the learning rate search results\n",
    "        lr_finder.plot()\n",
    "    \"\"\"\n",
    "    def __init__(self, model, criterion, optimizer, device):\n",
    "        \"\"\"\n",
    "        Initialize the LearningRateFinder.\n",
    "\n",
    "        Args:\n",
    "            model (torch.nn.Module): The neural network model to train.\n",
    "            criterion: The loss function used for training.\n",
    "            optimizer: The optimizer used for updating model parameters.\n",
    "            device (torch.device): The device (CPU or GPU) on which to perform training.\n",
    "        \"\"\"\n",
    "        self.model = model\n",
    "        self.criterion = criterion\n",
    "        self.optimizer = optimizer\n",
    "        self.device = device\n",
    "        self.history = {\"lr\": [], \"loss\": []}\n",
    "        self.best_loss = None\n",
    "        \n",
    "        self.model.to(device)\n",
    "        \n",
    "    def find_lr(\n",
    "        self, train_dataloader, \n",
    "        start_lr=None, end_lr=10, \n",
    "        num_iter=100, step_mode=\"exp\", \n",
    "        smooth_f=0.05, diverge_th=5,\n",
    "        early_stop_patience=None\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Search for the optimal learning rate within the specified range.\n",
    "\n",
    "        Args:\n",
    "            train_dataloader: The dataloader providing training batches.\n",
    "            start_lr (float or List[float], optional): The starting learning rate(s).\n",
    "                If None, uses the learning rate(s) from the optimizer. Default is None.\n",
    "            end_lr (float): The ending learning rate for the search.\n",
    "            num_iter (int): The number of iterations for the search.\n",
    "            step_mode (str, optional): The mode for adjusting learning rates (\"exp\" or \"linear\").\n",
    "                Default is \"exp\".\n",
    "            smooth_f (float, optional): The smoothing factor for loss values. Default is 0.05.\n",
    "            diverge_th (float, optional): The threshold for loss divergence. Default is 5.\n",
    "            early_stop_patience (int, optional): The patience for early stopping. Default is None.\n",
    "\n",
    "        Returns:\n",
    "            None\n",
    "        \"\"\"\n",
    "        # Initialize variables to track learning rate search history\n",
    "        self.history = {\"lr\": [], \"loss\": []}\n",
    "        self.best_loss = None # Best loss during the learning rate search\n",
    "        early_stop_counter = 0 # Counter for early stopping criteria\n",
    "\n",
    "        # Set initial learning rate if provided\n",
    "        if start_lr:\n",
    "            self.set_learning_rate(start_lr)\n",
    "\n",
    "        # Choose a learning rate scheduler based on the specified step_mode\n",
    "        if step_mode == \"exp\":\n",
    "            lr_schedule = ExponentialLR(self.optimizer, end_lr=end_lr, num_iter=num_iter) \n",
    "        elif step_mode == \"linear\":\n",
    "            lr_schedule = LinearLR(self.optimizer, end_lr=end_lr, num_iter=num_iter)\n",
    "\n",
    "        # Create an iterator for the training dataloader\n",
    "        train_iter = TrainDataLoaderIter(train_dataloader)\n",
    "\n",
    "        # Iterate over the specified number of iterations to find the optimal learning rate\n",
    "        for iteration in tqdm(range(num_iter)):\n",
    "            # Train a batch and obtain the loss\n",
    "            loss = self._train_batch(\n",
    "                train_iter\n",
    "            )\n",
    "            # Record the current learning rate in the history\n",
    "            self.history[\"lr\"].append(lr_schedule.get_lr()[0])\n",
    "            lr_schedule.step() # Update the learning rate according to the scheduler\n",
    "\n",
    "            # Update the best_loss if this is the first iteration, or if a lower loss is encountered\n",
    "            if iteration == 0:\n",
    "                self.best_loss = loss\n",
    "            else:\n",
    "                # Apply smoothing to the loss if specified\n",
    "                if smooth_f > 0:\n",
    "                    loss = smooth_f * loss + (1 - smooth_f) * self.history[\"loss\"][-1]\n",
    "\n",
    "                # Update best_loss and early stopping counter if a lower loss is not encountered\n",
    "                if loss < self.best_loss:\n",
    "                    self.best_loss = loss\n",
    "                else:\n",
    "                    # Check for early stopping based on patience\n",
    "                    early_stop_counter += 1\n",
    "                    if early_stop_patience and early_stop_counter > early_stop_patience:\n",
    "                        print(f\"Early stopping: Loss has not improved for {early_stop_patience} iterations.\")\n",
    "                        break\n",
    "\n",
    "            # Record the current loss in the history\n",
    "            self.history[\"loss\"].append(loss)\n",
    "            # Check for divergence in loss and stop the search if criteria are met\n",
    "            if loss > diverge_th * self.best_loss:\n",
    "                print(\"Stopping early, the loss has diverged.\")\n",
    "                break\n",
    "\n",
    "        # Print a message indicating the completion of the learning rate search\n",
    "        print(\"Learning rate search finished.\")\n",
    "        \n",
    "    def _train_batch(self, train_iter):\n",
    "        \"\"\"\n",
    "        Train a single batch and update the model parameters.\n",
    "\n",
    "        Args:\n",
    "            train_iter: Iterator providing training batches.\n",
    "\n",
    "        Returns:\n",
    "            float: The loss value for the current batch.\n",
    "        \"\"\"\n",
    "        # Set the model in training mode and zero the gradients\n",
    "        self.model.train()\n",
    "        self.optimizer.zero_grad()\n",
    "\n",
    "        # Retrieve inputs and labels for the current batch\n",
    "        inputs, labels = next(train_iter)\n",
    "        inputs, labels = inputs.to(self.device), labels.to(self.device)\n",
    "\n",
    "        # Forward pass: compute model predictions and calculate the loss\n",
    "        outputs = self.model(inputs)\n",
    "        loss = self.criterion(outputs, labels)\n",
    "\n",
    "        # Backward pass: compute gradients and update model parameters\n",
    "        loss.backward()\n",
    "        self.optimizer.step()\n",
    "        \n",
    "        # Return the loss value as a float\n",
    "        return loss.item()\n",
    "    \n",
    "    def set_learning_rate(self, start_lrs):\n",
    "        \"\"\"\n",
    "        Set the learning rate(s) for the optimizer.\n",
    "\n",
    "        Args:\n",
    "            start_lrs (float or List[float]): The initial learning rate(s) to be set.\n",
    "\n",
    "        Returns:\n",
    "            None\n",
    "        \"\"\"\n",
    "        # Ensure start_lrs is a list to handle single and multiple parameter groups\n",
    "        if not isinstance(start_lrs, list):\n",
    "            start_lrs = [start_lrs] * len(self.optimizer.param_groups)\n",
    "\n",
    "        # Update the learning rate for each parameter group\n",
    "        for param_group, start_lr in zip(self.optimizer.param_groups, start_lrs):\n",
    "            param_group['lr'] = start_lr\n",
    "            \n",
    "    def plot(self, skip_start = 10, skip_end=5, log_lr=True, show_lr=None):\n",
    "        \"\"\"\n",
    "        Plot the learning rate search results and suggest an optimal learning rate.\n",
    "\n",
    "        Args:\n",
    "            skip_start (int, optional): Number of initial iterations to skip in the plot. Default is 10.\n",
    "            skip_end (int, optional): Number of final iterations to skip in the plot. Default is 5.\n",
    "            log_lr (bool, optional): Whether to use a logarithmic scale for the learning rate axis.\n",
    "                Default is True.\n",
    "            show_lr (float, optional): The learning rate to highlight on the plot.\n",
    "\n",
    "        Returns:\n",
    "            tuple: A tuple containing the plot axis and the suggested optimal learning rate.\n",
    "        \"\"\"\n",
    "        lrs = self.history[\"lr\"]\n",
    "        losses = self.history[\"loss\"]\n",
    "        if skip_end == 0:\n",
    "            lrs = lrs[skip_start:]\n",
    "            losses = losses[skip_start:]\n",
    "        else:\n",
    "            lrs = lrs[skip_start:-skip_end]\n",
    "            losses = losses[skip_start:-skip_end]\n",
    "            \n",
    "        fig, ax = plt.subplots()    \n",
    "        ax.plot(lrs, losses)\n",
    "        \n",
    "        print(\"LR suggestion: steepest gradient\")\n",
    "        min_grad_idx = (np.gradient(np.array(losses))).argmin()\n",
    "        print(\"Suggested LR: {:.2E}\".format(lrs[min_grad_idx]))\n",
    "        ax.scatter(\n",
    "            lrs[min_grad_idx],\n",
    "            losses[min_grad_idx],\n",
    "            s=75,\n",
    "            marker=\"o\",\n",
    "            color=\"red\",\n",
    "            zorder=3,\n",
    "            label=\"steepest gradient\",\n",
    "        )\n",
    "        ax.legend()\n",
    "        \n",
    "        if log_lr:\n",
    "            ax.set_xscale(\"log\")\n",
    "        ax.set_xlabel(\"Learning rate\")\n",
    "        ax.set_ylabel(\"Loss\")\n",
    "        \n",
    "        if show_lr is not None:\n",
    "            ax.axvline(x=show_lr, color=\"red\")\n",
    "            \n",
    "        plt.show()\n",
    "            \n",
    "        return ax, lrs[min_grad_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2e6a0e9d-798f-4adc-b110-ef203922bb6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# lr_finder = LearningRateFinder(model, criterion, optimizer, device)\n",
    "# lr_finder.find_lr(train_dataloader, num_iter=500, start_lr=1e-3)\n",
    "# lr_finder.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6297eb60-ff83-4f9a-abbb-47919289ecc4",
   "metadata": {},
   "source": [
    "#### Train\n",
    "\n",
    "##### Gradient Penalty\n",
    "Its primary purpose is to encourage smoothness in the learned representations of the model by penalizing sharp changes or oscillations in the model's gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "34cd9acf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    model.train()\n",
    "    total_loss = 0.0\n",
    "\n",
    "    with tqdm(total=len(train_dataloader), desc=f\"Epoch {epoch}/Training\", unit=\"batch\") as pbar:\n",
    "        for i, (images, labels) in enumerate(train_dataloader):\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            output = model(images)\n",
    "            loss = criterion(output, labels)\n",
    "\n",
    "            # Computes the penalty term and adds it to the loss\n",
    "            grad_params = torch.autograd.grad(outputs=loss, inputs=model.parameters(), create_graph=True)\n",
    "            gradient_penalty = 0\n",
    "            for grad in grad_params:\n",
    "                gradient_penalty += grad.pow(2).sum()\n",
    "            gradient_penalty = gradient_penalty.sqrt()\n",
    "\n",
    "            loss = loss + gradient_penalty\n",
    "            loss.backward()\n",
    "\n",
    "            # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)\n",
    "            adaptive_clip_grad(model.parameters())\n",
    "            \n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            pbar.update(1)\n",
    "            pbar.set_postfix(loss=total_loss / (i+1), lr=f\"{optimizer.param_groups[0]['lr']:.5f}\")\n",
    "\n",
    "    avg_loss = total_loss / len(train_dataloader.dataset)\n",
    "    print(f\"[Train][Epoch {epoch}] Average Loss: {avg_loss:.5f}, Updated Learning Rate: {optimizer.param_groups[0]['lr']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a217af40-7a4e-4a1e-8cfc-133fc220090f",
   "metadata": {},
   "source": [
    "#### Eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "de2f38e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(epoch):\n",
    "    model.eval()\n",
    "    total_correct = 0\n",
    "    total_loss = 0.0\n",
    "\n",
    "    with tqdm(total=len(test_dataloader), desc=f\"Epoch {epoch}/Testing\", unit=\"batch\") as pbar: \n",
    "        with torch.no_grad():\n",
    "            for i, (images, labels) in enumerate(test_dataloader):\n",
    "                images, labels = images.to(device), labels.to(device)\n",
    "                \n",
    "                output = model(images)                \n",
    "                loss = criterion(output, labels)\n",
    "\n",
    "                total_loss += loss.item()\n",
    "                pbar.update(1)\n",
    "                pbar.set_postfix(loss=total_loss / (i+1))\n",
    "                \n",
    "                pred = output.detach().max(1)[1]\n",
    "                total_correct += pred.eq(labels.view_as(pred)).sum()\n",
    "            \n",
    "    avg_loss = total_loss / len(test_dataloader.dataset)\n",
    "    accuracy = total_correct / len(test_dataloader.dataset)\n",
    "    \n",
    "    print(f\"[Test][Epoch {epoch}] Loss: {avg_loss:.5f}, Accuracy: {accuracy:3f}\")\n",
    "    return avg_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e589a1d-848f-4045-813c-662ef287e348",
   "metadata": {},
   "source": [
    "#### Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "0c3cd024",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run():\n",
    "    for epoch in range(EPOCHS):\n",
    "        train(epoch)\n",
    "        val_loss = test(epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "04e628b2-39af-4e82-8d13-e94a3f4ffd86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ab479a406eff41638671bb46371bae9e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 0/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 0] Average Loss: 0.00686, Updated Learning Rate: 0.0006708978\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64fbcc46eacd4d288744f1e9dc9d4de2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 0/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 0] Loss: 0.00322, Accuracy: 0.608100\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "018748745fa94e86b2baa1bf623306ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 1/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 1] Average Loss: 0.00442, Updated Learning Rate: 0.00134\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1ac18797e0634bf7a77cb6a72252e48c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 1/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 1] Loss: 0.00227, Accuracy: 0.808600\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d665fccc83b2490196c9408eb26e8a2d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 2/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 2] Average Loss: 0.00354, Updated Learning Rate: 0.0013298211945181795\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4bc836f48abb418fa8a8b713741722b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 2/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 2] Loss: 0.00185, Accuracy: 0.882100\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ab732b5b7214882865fc71f3a06a407",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 3/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 3] Average Loss: 0.00312, Updated Learning Rate: 0.0012995940559265588\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fadb8be1d7f346ad97b621324a34ff27",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 3/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 3] Loss: 0.00168, Accuracy: 0.914600\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9fbeab63b8240ea96b6227009703092",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 4/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 4] Average Loss: 0.00288, Updated Learning Rate: 0.0012502370205355739\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d2d0095403b431596f7141f85ecc2d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 4/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 4] Loss: 0.00157, Accuracy: 0.934600\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f6773e8540864b1283cef4b5851d7995",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 5/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 5] Average Loss: 0.00273, Updated Learning Rate: 0.0011832497768897153\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d0d006a309b4e7db970ed5f74e24e9c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 5/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 5] Loss: 0.00153, Accuracy: 0.938200\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ece8e284b424540bed3d0c40b50a4c7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 6/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 6] Average Loss: 0.00262, Updated Learning Rate: 0.0011006676984899814\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "df502468eca14dddb30e172655230d00",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 6/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 6] Loss: 0.00148, Accuracy: 0.947500\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "13ef2bcfc8544f08b99ded589592fc4d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 7/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 7] Average Loss: 0.00254, Updated Learning Rate: 0.001005\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfe7bdb6230848bca24db541b26c2e60",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 7/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 7] Loss: 0.00146, Accuracy: 0.951000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6a9554ca8d9b4877a1e0259d65a083f0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 8/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 8] Average Loss: 0.00246, Updated Learning Rate: 0.0008991534960281982\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c41b2218df3c4efa870be3b036898614",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 8/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 8] Loss: 0.00144, Accuracy: 0.950800\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34564ce360484b9ca6d850e78ca3d273",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 9/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 9] Average Loss: 0.00241, Updated Learning Rate: 0.0007863442790368435\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f89db048eb21455b964914f2475e9d8e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 9/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 9] Loss: 0.00143, Accuracy: 0.953500\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1f6b2bcb1e534ee9b0c1d696bb08235f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 10/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 10] Average Loss: 0.00238, Updated Learning Rate: 0.00067\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d2e75e92aea64ef8aed7e3a6084eda4f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 10/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 10] Loss: 0.00142, Accuracy: 0.956800\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "911d8772954840b196f0882dc117597d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 11/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 11] Average Loss: 0.00235, Updated Learning Rate: 0.0005536557209631567\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a42481f9d03c4ab59112ce0168de3ecf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 11/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 11] Loss: 0.00141, Accuracy: 0.957000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac5b5b05a2304604a5f4a76f7abbfa35",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 12/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 12] Average Loss: 0.00232, Updated Learning Rate: 0.000440846503971802\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e96a1817c2240b7bf735bc4c590c5ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 12/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 12] Loss: 0.00141, Accuracy: 0.954500\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aaf4731ab9e9405fb004d62c7a002df8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 13/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 13] Average Loss: 0.00231, Updated Learning Rate: 0.00033499999999999985\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c4788a793576485d8e378c2c8d2d4338",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 13/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 13] Loss: 0.00140, Accuracy: 0.956500\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ccdc9dfa38594f3193ede8443f347ab3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 14/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 14] Average Loss: 0.00228, Updated Learning Rate: 0.00023933230151001864\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d1645d5991be41f8b956fab928d0e061",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 14/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 14] Loss: 0.00139, Accuracy: 0.959500\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e9e3d5ca4d2449eba0db2c97b98c620",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 15/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 15] Average Loss: 0.00227, Updated Learning Rate: 0.00015675022311028482\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "900189f172dd40f1b4cb3aeccc447811",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 15/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 15] Loss: 0.00140, Accuracy: 0.957700\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3d849a10c00b450bb20293e6d8c58a3d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 16/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 16] Average Loss: 0.00226, Updated Learning Rate: 8.976297946442607e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e444400d4e9432ba1efce1c710c855d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 16/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 16] Loss: 0.00139, Accuracy: 0.959000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ae8bff4cc1b4c9e818e4dcb4826215d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 17/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 17] Average Loss: 0.00228, Updated Learning Rate: 4.040594407344143e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c20b7d1a5ffc43ffa4d85117b50902a1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 17/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 17] Loss: 0.00139, Accuracy: 0.958800\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a8cb33ceb38c4d689a99cf57e7ecc019",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 18/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 18] Average Loss: 0.00226, Updated Learning Rate: 1.0178805481820626e-05\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dceaa0779a8c455688955eb5605622ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 18/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 18] Loss: 0.00139, Accuracy: 0.958500\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c8005fcd65de4be4913009e69544118d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 19/Training:   0%|          | 0/118 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 19] Average Loss: 0.00225, Updated Learning Rate: 0.0\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c26c54ffa424784a76a1264a3adf14b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 19/Testing:   0%|          | 0/20 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 19] Loss: 0.00139, Accuracy: 0.957900\n"
     ]
    }
   ],
   "source": [
    "run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7ee87b2-3663-4643-aeaf-05b7be98dd42",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b75b79a8-8176-4788-801c-7f04b2bce909",
   "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
}