08_Hyperparameter_Optimization/24_2_Hyperparameter_Tuning_With_optuna.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 time\n",
    "import math\n",
    "import datetime\n",
    "from collections import OrderedDict\n",
    "from tqdm.auto import tqdm\n",
    "from pathlib import Path\n",
    "from thop import profile\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 mlflow\n",
    "import random\n",
    "import numpy as np\n",
    "import psutil\n",
    "import subprocess\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": "13abf35c-ae3c-4ff4-8c4f-c5bbb27e87bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.cuda.memory._record_memory_history()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1db77c57-1975-4f86-9188-59fcc8203174",
   "metadata": {},
   "outputs": [],
   "source": [
    "def increment_path(path, exist_ok=False, sep='', mkdir=True):\n",
    "    path = Path(path)\n",
    "    if path.exists() and not exist_ok:\n",
    "        path, suffix = (path.with_suffix(''), path.suffix) if path.is_file() else (path, '')\n",
    "        for n in range(2, 9999):\n",
    "            p = f'{path}{sep}{n}{suffix}'\n",
    "            if not os.path.exists(p):\n",
    "                break\n",
    "        path = Path(p)\n",
    "    if mkdir:\n",
    "        path.mkdir(parents=True, exist_ok=True)\n",
    "    return path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "94032296-1658-4b03-bb63-5c2283e00ce5",
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPERIMENT_NAME = \"exp\"\n",
    "FILE = Path(os.getcwd()).resolve()\n",
    "SAVE_DIR = increment_path(FILE / \"artifacts\" / f\"{EXPERIMENT_NAME}\", exist_ok=False)\n",
    "(SAVE_DIR / 'checkpoints').mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "15c8f025-3784-446b-b6a6-f9570e0bec97",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/03/11 19:37:31 INFO mlflow.tracking.fluent: Experiment with name 'exp13' does not exist. Creating a new experiment.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Experiment: artifact_location='file:///C:/Users/Hp/Desktop/mldev/2024/Git/mldevworld/Deep-Learning-Essentials/mlruns/13', creation_time=1710166051934, experiment_id='13', last_update_time=1710166051934, lifecycle_stage='active', name='exp13', tags={}>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "local_tracking_uri = f\"sqlite:///mlflow/mlflow.db\"\n",
    "mlflow.set_tracking_uri(local_tracking_uri)\n",
    "mlflow.set_experiment(SAVE_DIR.stem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": 8,
   "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": 9,
   "id": "372c5b4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define hyperparameters\n",
    "N_CLASSES = 10\n",
    "EPOCHS = 10\n",
    "BATCH_SIZE = 1609\n",
    "DROPOUT_PROB = \t0.12931721102481436\n",
    "USE_BATCHNORM = False\n",
    "LABEL_SMOOTHING = 8.83890730090731e-07\n",
    "LR = 0.028147621640854794\n",
    "WARMUP_PROPORTION = 0.04584907346713653\n",
    "WEIGHT_DECAY = 2.1415275778699216e-06\n",
    "ACCUMULATION_STEPS = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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": 11,
   "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": 12,
   "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": 13,
   "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": 14,
   "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": 15,
   "id": "31f61bca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1          [1609, 6, 28, 28]             156\n",
      "          Identity-2          [1609, 6, 28, 28]               0\n",
      "              ReLU-3          [1609, 6, 28, 28]               0\n",
      "         MaxPool2d-4          [1609, 6, 14, 14]               0\n",
      "         Dropout2d-5          [1609, 6, 14, 14]               0\n",
      "            Conv2d-6         [1609, 16, 10, 10]           2,416\n",
      "          Identity-7         [1609, 16, 10, 10]               0\n",
      "              ReLU-8         [1609, 16, 10, 10]               0\n",
      "         MaxPool2d-9           [1609, 16, 5, 5]               0\n",
      "        Dropout2d-10           [1609, 16, 5, 5]               0\n",
      "           Conv2d-11          [1609, 120, 1, 1]          48,120\n",
      "         Identity-12          [1609, 120, 1, 1]               0\n",
      "             ReLU-13          [1609, 120, 1, 1]               0\n",
      "        Dropout2d-14          [1609, 120, 1, 1]               0\n",
      "           Linear-15                 [1609, 84]          10,164\n",
      "             ReLU-16                 [1609, 84]               0\n",
      "          Dropout-17                 [1609, 84]               0\n",
      "           Linear-18                 [1609, 10]             850\n",
      "================================================================\n",
      "Total params: 61,706\n",
      "Trainable params: 61,706\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 6.29\n",
      "Forward/backward pass size (MB): 279.96\n",
      "Params size (MB): 0.24\n",
      "Estimated Total Size (MB): 286.48\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": 16,
   "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": 17,
   "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": 18,
   "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": 19,
   "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": 20,
   "id": "b45993ef-a7a2-4154-9eeb-d10e5e003e4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "len_dataloader = len(train_dataloader)\n",
    "num_update_steps_per_epoch = max(len_dataloader // ACCUMULATION_STEPS, 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": 21,
   "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": "465ce651-83e8-4526-a9fc-bf799b70446b",
   "metadata": {},
   "source": [
    "#### Logging Model Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e8aa3340-60e2-4583-8867-d197fb467d33",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_flops(input_size):\n",
    "    input_tensor = torch.randn(*input_size).to(device)\n",
    "    flops, params = profile(model, inputs=(input_tensor,))\n",
    "    return flops\n",
    "\n",
    "def calculate_disk_size():\n",
    "    # Calculate model disk size in bytes\n",
    "    disk_size_bytes = sum(p.numel() * p.element_size() for p in model.parameters())\n",
    "    # Convert bytes to megabytes\n",
    "    disk_size_mb = disk_size_bytes / (1024 * 1024)\n",
    "    return disk_size_mb\n",
    "\n",
    "def get_model_metrics():\n",
    "    num_params = sum(p.numel() for p in model.parameters())\n",
    "    num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "    # Set an example input size (adjust based on your actual input size)\n",
    "    input_size = tuple(next(iter(train_dataloader))[0][0].unsqueeze(0).shape)\n",
    "\n",
    "    flops = calculate_flops(input_size)\n",
    "    disk_size = calculate_disk_size()\n",
    "\n",
    "    return {\n",
    "        \"num_params\": num_params,\n",
    "        \"num_trainable_params\": num_trainable_params,\n",
    "        \"flops\": flops,\n",
    "        \"disk_size\": disk_size\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0be3441b-14b2-4251-9bc4-b61b60263606",
   "metadata": {},
   "source": [
    "#### Logging GPU Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "02adfbae-1213-4960-b91a-0abf1079256a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_gpu_usage():\n",
    "    try:\n",
    "        # Run nvidia-smi command to get GPU utilization information\n",
    "        result = subprocess.run([\"nvidia-smi\", \"--query-gpu=utilization.gpu\", \"--format=csv,nounits,noheader\"], stdout=subprocess.PIPE, text=True)\n",
    "        gpu_usage = int(result.stdout.strip())\n",
    "        return gpu_usage\n",
    "    except Exception as e:\n",
    "        print(f\"Error getting GPU usage: {e}\")\n",
    "        return None"
   ]
  },
  {
   "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.\n",
    "\n",
    "##### Gradient Accumulation\n",
    "Gradient accumulation adds gradients over an effective batch of size `batch_per_iter * iters_to_accumulate` (`* num_procs` if distributed). The scale should be calibrated for the effective batch, which means inf/NaN checking, step skipping if inf/NaN grads are found, and scale updates should occur at effective-batch granularity. Also, grads should remain scaled, and the scale factor should remain constant, while grads for a given effective batch are accumulated. If grads are unscaled (or the scale factor changes) before accumulation is complete, the next backward pass will add scaled grads to unscaled grads (or grads scaled by a different factor) after which it’s impossible to recover the accumulated unscaled grads `step` must apply.\n",
    "\n",
    "##### Automatic Mixed Precision Training\n",
    "Automatic Mixed Precision (AMP) training is a technique used in deep learning to accelerate training by using a combination of lower-precision and higher-precision numerical representations. The primary idea behind AMP is to use lower-precision data types (such as float16) for some parts of the model computation, while maintaining higher precision (such as float32) for critical numerical stability aspects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "34cd9acf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    model.train()\n",
    "    total_loss = 0.0\n",
    "    \n",
    "    # Create a GradScaler once at the beginning of training.\n",
    "    scaler = torch.cuda.amp.GradScaler()\n",
    "    \n",
    "    optimizer.zero_grad()\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",
    "            # Runs the forward pass with autocasting.\n",
    "            with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=True):\n",
    "                output = model(images)\n",
    "                loss = criterion(output, labels)\n",
    "\n",
    "            # # Scales the loss for autograd.grad's backward pass, producing scaled_grad_params\n",
    "            # scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss), inputs=model.parameters(), create_graph=True)\n",
    "    \n",
    "            # # Creates unscaled grad_params before computing the penalty. scaled_grad_params are\n",
    "            # # not owned by any optimizer, so ordinary division is used instead of scaler.unscale_:\n",
    "            # inv_scale = 1. / scaler.get_scale()\n",
    "            # grad_params = [p * inv_scale for p in scaled_grad_params]\n",
    "    \n",
    "            # # Computes the penalty term and adds it to the loss\n",
    "            # with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=True):\n",
    "            #     grad_norm = 0\n",
    "            #     for grad in grad_params:\n",
    "            #         grad_norm += grad.pow(2).sum()\n",
    "            #     grad_norm = grad_norm.sqrt()\n",
    "            #     loss = loss + grad_norm\n",
    "\n",
    "            # Scale the loss by accumulation steps\n",
    "            if ACCUMULATION_STEPS > 1:\n",
    "                with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=True):\n",
    "                    loss = loss / ACCUMULATION_STEPS\n",
    "                \n",
    "            # Applies scaling to the backward call as usual.\n",
    "            # Accumulates leaf gradients that are correctly scaled.\n",
    "            scaler.scale(loss).backward()\n",
    "\n",
    "            if (i + 1) % ACCUMULATION_STEPS == 0:\n",
    "                # Unscales the gradients of optimizer's assigned params in-place\n",
    "                scaler.unscale_(optimizer)\n",
    "                # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)\n",
    "                adaptive_clip_grad(model.parameters())\n",
    "\n",
    "                # optimizer's gradients are already unscaled, so scaler.step does not unscale them,\n",
    "                # although it still skips optimizer.step() if the gradients contain infs or NaNs.\n",
    "                scaler.step(optimizer)\n",
    "\n",
    "                mlflow.log_metric(\"lr\", optimizer.param_groups[0][\"lr\"])\n",
    "                scheduler.step()\n",
    "\n",
    "                # Updates the scale for next iteration.\n",
    "                scaler.update()\n",
    "                \n",
    "                for param in model.parameters():\n",
    "                    param.grad = None\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",
    "    mlflow.log_metric(\"train_loss\", avg_loss)\n",
    "        \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": 25,
   "id": "65850723-b84a-4915-8429-644a1805ba97",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BestLossTracker:\n",
    "    def __init__(self):\n",
    "        self.best_loss = float('inf')\n",
    "    \n",
    "    def update_best_loss(self, avg_loss):\n",
    "        if avg_loss < self.best_loss:\n",
    "            self.best_loss = avg_loss\n",
    "            return True\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c2ca1f2b-086a-4160-a0a6-f0abb7c43100",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EarlyStopping:\n",
    "    def __init__(self, patience=5, delta=0.0):\n",
    "        self.patience = patience\n",
    "        self.delta = delta\n",
    "        self.counter = 0\n",
    "        self.best_loss = float('inf')\n",
    "        self.early_stop = False\n",
    "\n",
    "    def update_best_loss(self, val_loss):\n",
    "        if val_loss < self.best_loss - self.delta:\n",
    "            self.best_loss = val_loss\n",
    "            self.counter = 0\n",
    "        else:\n",
    "            self.counter += 1\n",
    "            if self.counter >= self.patience:\n",
    "                self.early_stop = True\n",
    "\n",
    "        return not self.early_stop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "de2f38e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(epoch, best_loss_tracker):\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",
    "    mlflow.log_metric(\"val_loss\", avg_loss)\n",
    "    mlflow.log_metric(\"val_accuracy\", accuracy)\n",
    "    \n",
    "    if best_loss_tracker.update_best_loss(avg_loss):\n",
    "        checkpoint_name = f\"checkpoint_best_epoch_{epoch}.pth\"\n",
    "        checkpoint_path = str(SAVE_DIR / \"checkpoints\" / checkpoint_name)\n",
    "        torch.save(model.state_dict(), checkpoint_path)\n",
    "        print(f\"[Test][Epoch {epoch}] New best model found! Saving checkpoint. Loss: {avg_loss:.5f}, Accuracy: {accuracy:.3f}\")\n",
    "    else:\n",
    "        print(f\"[Test][Epoch {epoch}] Loss: {avg_loss:.5f}, Accuracy: {accuracy:.3f}\")\n",
    "    \n",
    "    return avg_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e589a1d-848f-4045-813c-662ef287e348",
   "metadata": {},
   "source": [
    "#### Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "0c3cd024",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run():\n",
    "    with mlflow.start_run(run_name=\"Understanding_CUDA_Memory_Usage\"):\n",
    "        mlflow.log_param(\"torch_version\", torch.__version__)\n",
    "        mlflow.log_param(\"mlflow_version\", mlflow.__version__)\n",
    "\n",
    "        hyper_params = {\n",
    "            \"accumulation_steps\": ACCUMULATION_STEPS, \n",
    "            \"batch_size\": BATCH_SIZE, \n",
    "            \"dropout\":DROPOUT_PROB,\n",
    "            \"epochs\": EPOCHS, \n",
    "            \"batchnorm\": USE_BATCHNORM,\n",
    "            \"label_smoothing\": LABEL_SMOOTHING,\n",
    "            \"start_lr\": LR,\n",
    "            \"warmup_proortion\": WARMUP_PROPORTION,\n",
    "            \"weight_decay\": WEIGHT_DECAY\n",
    "        }\n",
    "\n",
    "        mlflow.log_params(hyper_params)\n",
    "        mlflow.set_tags({\"model\": \"lenet\", \"dataset\": \"mnist\"})\n",
    "        mlflow.pytorch.log_model(model, \"model\")\n",
    "        mlflow.log_param(\"start_time\", datetime.datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\"))\n",
    "\n",
    "        model_metrics = get_model_metrics()\n",
    "        mlflow.log_param(\"num_params\", model_metrics['num_params'])\n",
    "        mlflow.log_param(\"num_trainable_params\", model_metrics['num_trainable_params'])\n",
    "        mlflow.log_param(\"flops\", model_metrics['flops'])\n",
    "        mlflow.log_param(\"disk_size\", model_metrics['disk_size'])\n",
    "\n",
    "        best_loss_tracker = BestLossTracker()\n",
    "        early_stopping = EarlyStopping(patience=3, delta=1e-6)\n",
    "\n",
    "        start_time = time.time()\n",
    "        for epoch in range(EPOCHS):\n",
    "            train(epoch)\n",
    "            val_loss = test(epoch, best_loss_tracker)\n",
    "            \n",
    "            if get_gpu_usage() is not None:\n",
    "                mlflow.log_metric(\"gpu\", get_gpu_usage(), step=epoch)\n",
    "\n",
    "            if not early_stopping.update_best_loss(val_loss):\n",
    "                print(\"Early stopping triggered.\")\n",
    "                break\n",
    "\n",
    "\n",
    "        total_time = time.time() - start_time\n",
    "        avg_epoch_time = total_time / EPOCHS\n",
    "        avg_batch_time = avg_epoch_time / len(train_dataloader)\n",
    "\n",
    "        mlflow.log_metric(\"total_training_time\", total_time)\n",
    "        mlflow.log_metric(\"average_epoch_time\", avg_epoch_time)\n",
    "        mlflow.log_metric(\"average_batch_time\", avg_batch_time)\n",
    "        mlflow.log_param(\"end_time\", datetime.datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\"))\n",
    "\n",
    "# torch.cuda.memory._dump_snapshot(\"snapshot.pickle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "04e628b2-39af-4e82-8d13-e94a3f4ffd86",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/03/11 19:37:35 WARNING mlflow.utils.requirements_utils: Found torch version (2.2.0+cu118) contains a local version label (+cu118). MLflow logged a pip requirement for this package as 'torch==2.2.0' without the local version label to make it installable from PyPI. To specify pip requirements containing local version labels, please use `conda_env` or `pip_requirements`.\n",
      "2024/03/11 19:37:38 WARNING mlflow.utils.requirements_utils: Found torch version (2.2.0+cu118) contains a local version label (+cu118). MLflow logged a pip requirement for this package as 'torch==2.2.0' without the local version label to make it installable from PyPI. To specify pip requirements containing local version labels, please use `conda_env` or `pip_requirements`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.\n",
      "[INFO] Register zero_ops() for <class 'torch.nn.modules.activation.ReLU'>.\n",
      "[INFO] Register zero_ops() for <class 'torch.nn.modules.pooling.MaxPool2d'>.\n",
      "[INFO] Register zero_ops() for <class 'torch.nn.modules.container.Sequential'>.\n",
      "[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.\n",
      "[INFO] Register zero_ops() for <class 'torch.nn.modules.dropout.Dropout'>.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be0e3155e8c748f78c9aab0905183a24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 0/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 0] Average Loss: 0.00090, Updated Learning Rate: 0.027936158958914464\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f008bece9cea4da2878a74490d3a82af",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 0/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 0] New best model found! Saving checkpoint. Loss: 0.00016, Accuracy: 0.930\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "442a82259f774ddb8bb91fed5f08701a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 1/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 1] Average Loss: 0.00036, Updated Learning Rate: 0.026402075025726137\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f8c1276daf7f48b9a238e1a1768ba873",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 1/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 1] New best model found! Saving checkpoint. Loss: 0.00014, Accuracy: 0.937\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "04ceb1eae0ff44979643085556c8a0ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 2/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 2] Average Loss: 0.00028, Updated Learning Rate: 0.023539336241444258\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0a1bc7ab5652479e9816d646dc658856",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 2/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 2] New best model found! Saving checkpoint. Loss: 0.00008, Accuracy: 0.963\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7b15aecae74948369941bb7831b1966c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 3/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 3] Average Loss: 0.00023, Updated Learning Rate: 0.019656468746473504\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9dd66c79a3d94a0cbb444be5cdef502c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 3/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 3] New best model found! Saving checkpoint. Loss: 0.00007, Accuracy: 0.970\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d090394b1f1d48cdae3e042437bf2142",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 4/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 4] Average Loss: 0.00021, Updated Learning Rate: 0.015171941081675348\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1347cad9c3014762992a09e4bbf7ccc2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 4/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 4] New best model found! Saving checkpoint. Loss: 0.00007, Accuracy: 0.970\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fed82fcb37364227953c15456edcdabe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 5/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 5] Average Loss: 0.00019, Updated Learning Rate: 0.010569064550042183\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e78b7239c8624f49b1e1e6ee76cce07c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 5/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 5] New best model found! Saving checkpoint. Loss: 0.00007, Accuracy: 0.970\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c204790dd16496793cc284eb2e294a9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 6/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 6] Average Loss: 0.00017, Updated Learning Rate: 0.006343905274987439\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e30cf3301d442c793bea177744c7dd4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 6/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 6] New best model found! Saving checkpoint. Loss: 0.00006, Accuracy: 0.976\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "65afd388ec104bc688bcb688e049a8d3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 7/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 7] Average Loss: 0.00016, Updated Learning Rate: 0.0029518216325127794\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "562a2b3a3e4c40e0b3541f4eea0bbd35",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 7/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 7] New best model found! Saving checkpoint. Loss: 0.00005, Accuracy: 0.975\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fcc712f1ee4f40d191d79f8d556b933d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 8/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 8] Average Loss: 0.00015, Updated Learning Rate: 0.000758388882187354\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e12542114244f64a480f9a3d87096fa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 8/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 8] New best model found! Saving checkpoint. Loss: 0.00005, Accuracy: 0.978\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd27b2c7756d4506bbe79894fdd30195",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 9/Training:   0%|          | 0/38 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Train][Epoch 9] Average Loss: 0.00015, Updated Learning Rate: 0.0\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "28d4cee6ac664869b9290a0342dd1b75",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 9/Testing:   0%|          | 0/7 [00:00<?, ?batch/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Test][Epoch 9] Loss: 0.00005, Accuracy: 0.977\n"
     ]
    }
   ],
   "source": [
    "run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "555bbc83-0f0d-45c4-a0e3-e7ea6f441d8c",
   "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
}