docs/usercases/tracking-mljob-in-kata-containers/cifar10-rest50-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Code from https://www.learnpytorch.io/pytorch_2_intro, under MIT LICENSE.\n",
    "\n",
    "Thanks for the awesome tutorial!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://download.pytorch.org/whl/cu117\n",
      "Requirement already satisfied: torch in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (2.0.1+cu117)\n",
      "Requirement already satisfied: torchvision in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (0.15.2+cu117)\n",
      "Requirement already satisfied: torchaudio in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (2.0.2+cu117)\n",
      "Requirement already satisfied: filelock in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from torch) (3.9.0)\n",
      "Requirement already satisfied: typing-extensions in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from torch) (4.7.1)\n",
      "Requirement already satisfied: sympy in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from torch) (1.11.1)\n",
      "Requirement already satisfied: networkx in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from torch) (3.1)\n",
      "Requirement already satisfied: jinja2 in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from torch) (3.1.2)\n",
      "Requirement already satisfied: triton==2.0.0 in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from torch) (2.0.0)\n",
      "Requirement already satisfied: cmake in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from triton==2.0.0->torch) (3.25.0)\n",
      "Requirement already satisfied: lit in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from triton==2.0.0->torch) (15.0.7)\n",
      "Requirement already satisfied: numpy in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from torchvision) (1.24.1)\n",
      "Requirement already satisfied: requests in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from torchvision) (2.31.0)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from torchvision) (9.3.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from jinja2->torch) (2.1.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from requests->torchvision) (3.2.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from requests->torchvision) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from requests->torchvision) (2.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from requests->torchvision) (2023.7.22)\n",
      "Requirement already satisfied: mpmath>=0.19 in /home/wh1isper/miniconda3/envs/runtime/lib/python3.10/site-packages (from sympy->torch) (1.2.1)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install torch torchvision tqdm torchaudio --index-url https://download.pytorch.org/whl/cu117"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch version: 2.0.1+cu117\n",
      "TorchVision version: 0.15.2+cu117\n",
      "Using device: cuda\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "print(f\"PyTorch version: {torch.__version__}\")\n",
    "print(f\"TorchVision version: {torchvision.__version__}\")\n",
    "\n",
    "# Set the target device\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "print(f\"Using device: {device}\")\n",
    "GPU_SCORE = torch.cuda.get_device_capability()\n",
    "print(f\"GPU capability score: {GPU_SCORE}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://download.pytorch.org/models/resnet50-11ad3fa6.pth\" to /home/wh1isper/.cache/torch/hub/checkpoints/resnet50-11ad3fa6.pth\n",
      "100.0%\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total parameters of model: 25557032 (the more parameters, the more GPU memory the model will use, the more *relative* of a speedup you'll get)\n",
      "Model transforms:\n",
      "ImageClassification(\n",
      "    crop_size=[224]\n",
      "    resize_size=[232]\n",
      "    mean=[0.485, 0.456, 0.406]\n",
      "    std=[0.229, 0.224, 0.225]\n",
      "    interpolation=InterpolationMode.BILINEAR\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Create model weights and transforms\n",
    "model_weights = (\n",
    "    torchvision.models.ResNet50_Weights.IMAGENET1K_V2\n",
    ")  # <- use the latest weights (could also use .DEFAULT)\n",
    "transforms = model_weights.transforms()\n",
    "\n",
    "# Setup model\n",
    "model = torchvision.models.resnet50(weights=model_weights)\n",
    "\n",
    "# Count the number of parameters in the model\n",
    "total_params = sum(\n",
    "    param.numel()\n",
    "    for param in model.parameters()  # <- all params\n",
    "    # param.numel() for param in model.parameters() if param.requires_grad # <- only trainable params\n",
    ")\n",
    "\n",
    "print(\n",
    "    f\"Total parameters of model: {total_params} (the more parameters, the more GPU memory the model will use, the more *relative* of a speedup you'll get)\"\n",
    ")\n",
    "print(f\"Model transforms:\\n{transforms}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model(num_classes=10):\n",
    "    \"\"\"\n",
    "    Creates a ResNet50 model with the latest weights and transforms via torchvision.\n",
    "    \"\"\"\n",
    "    model_weights = torchvision.models.ResNet50_Weights.IMAGENET1K_V2\n",
    "    transforms = model_weights.transforms()\n",
    "    model = torchvision.models.resnet50(weights=model_weights)\n",
    "\n",
    "    # Adjust the number of output features in model to match the number of classes in the dataset\n",
    "    model.fc = torch.nn.Linear(in_features=2048, out_features=num_classes)\n",
    "    return model, transforms\n",
    "\n",
    "\n",
    "model, transforms = create_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total free GPU memory: 3.993 GB\n",
      "Total GPU memory: 4.235 GB\n"
     ]
    }
   ],
   "source": [
    "# Check available GPU memory and total GPU memory\n",
    "total_free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info()\n",
    "print(f\"Total free GPU memory: {round(total_free_gpu_memory * 1e-9, 3)} GB\")\n",
    "print(f\"Total GPU memory: {round(total_gpu_memory * 1e-9, 3)} GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU memory available is 3.993 GB, using batch size of 32 and image size 128\n"
     ]
    }
   ],
   "source": [
    "# Set batch size depending on amount of GPU memory\n",
    "total_free_gpu_memory_gb = round(total_free_gpu_memory * 1e-9, 3)\n",
    "if total_free_gpu_memory_gb >= 16:\n",
    "    BATCH_SIZE = 128  # Note: you could experiment with higher values here if you like.\n",
    "    IMAGE_SIZE = 224\n",
    "    print(\n",
    "        f\"GPU memory available is {total_free_gpu_memory_gb} GB, using batch size of {BATCH_SIZE} and image size {IMAGE_SIZE}\"\n",
    "    )\n",
    "else:\n",
    "    BATCH_SIZE = 32\n",
    "    IMAGE_SIZE = 128\n",
    "    print(\n",
    "        f\"GPU memory available is {total_free_gpu_memory_gb} GB, using batch size of {BATCH_SIZE} and image size {IMAGE_SIZE}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updated data transforms:\n",
      "ImageClassification(\n",
      "    crop_size=128\n",
      "    resize_size=128\n",
      "    mean=[0.485, 0.456, 0.406]\n",
      "    std=[0.229, 0.224, 0.225]\n",
      "    interpolation=InterpolationMode.BILINEAR\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "transforms.crop_size = IMAGE_SIZE\n",
    "transforms.resize_size = IMAGE_SIZE\n",
    "print(f\"Updated data transforms:\\n{transforms}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] Using GPU with score: (6, 1), TensorFloat32 (TF32) not available, to use it you need a GPU with score >= (8, 0)\n"
     ]
    }
   ],
   "source": [
    "if GPU_SCORE >= (8, 0):\n",
    "    print(\n",
    "        f\"[INFO] Using GPU with score: {GPU_SCORE}, enabling TensorFloat32 (TF32) computing (faster on new GPUs)\"\n",
    "    )\n",
    "    torch.backends.cuda.matmul.allow_tf32 = True\n",
    "else:\n",
    "    print(\n",
    "        f\"[INFO] Using GPU with score: {GPU_SCORE}, TensorFloat32 (TF32) not available, to use it you need a GPU with score >= (8, 0)\"\n",
    "    )\n",
    "    torch.backends.cuda.matmul.allow_tf32 = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./cifar-10-python.tar.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100.0%\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./cifar-10-python.tar.gz to .\n",
      "Files already downloaded and verified\n",
      "[INFO] Train dataset length: 50000\n",
      "[INFO] Test dataset length: 10000\n"
     ]
    }
   ],
   "source": [
    "# Create train and test datasets\n",
    "train_dataset = torchvision.datasets.CIFAR10(\n",
    "    root=\".\", train=True, download=True, transform=transforms\n",
    ")\n",
    "\n",
    "test_dataset = torchvision.datasets.CIFAR10(\n",
    "    root=\".\",\n",
    "    train=False,  # want the test split\n",
    "    download=True,\n",
    "    transform=transforms,\n",
    ")\n",
    "\n",
    "# Get the lengths of the datasets\n",
    "train_len = len(train_dataset)\n",
    "test_len = len(test_dataset)\n",
    "\n",
    "print(f\"[INFO] Train dataset length: {train_len}\")\n",
    "print(f\"[INFO] Test dataset length: {test_len}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train dataloader length: 1563 batches of size 32\n",
      "Test dataloader length: 313 batches of size 32\n",
      "Using number of workers: 8 (generally more workers means faster dataloading from CPU to GPU)\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Create DataLoaders\n",
    "import os\n",
    "\n",
    "NUM_WORKERS = os.cpu_count()  # <- use all available CPU cores (this number can be tweaked through experimentation but generally more workers means faster dataloading from CPU to GPU)\n",
    "\n",
    "train_dataloader = DataLoader(\n",
    "    dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS\n",
    ")\n",
    "\n",
    "test_dataloader = DataLoader(\n",
    "    dataset=test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS\n",
    ")\n",
    "\n",
    "# Print details\n",
    "print(f\"Train dataloader length: {len(train_dataloader)} batches of size {BATCH_SIZE}\")\n",
    "print(f\"Test dataloader length: {len(test_dataloader)} batches of size {BATCH_SIZE}\")\n",
    "print(\n",
    "    f\"Using number of workers: {NUM_WORKERS} (generally more workers means faster dataloading from CPU to GPU)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from tqdm.auto import tqdm\n",
    "from typing import Dict, List, Tuple\n",
    "\n",
    "\n",
    "def train_step(\n",
    "    epoch: int,\n",
    "    model: torch.nn.Module,\n",
    "    dataloader: torch.utils.data.DataLoader,\n",
    "    loss_fn: torch.nn.Module,\n",
    "    optimizer: torch.optim.Optimizer,\n",
    "    device: torch.device,\n",
    "    disable_progress_bar: bool = False,\n",
    ") -> Tuple[float, float]:\n",
    "    \"\"\"Trains a PyTorch model for a single epoch.\n",
    "\n",
    "    Turns a target PyTorch model to training mode and then\n",
    "    runs through all of the required training steps (forward\n",
    "    pass, loss calculation, optimizer step).\n",
    "\n",
    "    Args:\n",
    "      model: A PyTorch model to be trained.\n",
    "      dataloader: A DataLoader instance for the model to be trained on.\n",
    "      loss_fn: A PyTorch loss function to minimize.\n",
    "      optimizer: A PyTorch optimizer to help minimize the loss function.\n",
    "      device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "      A tuple of training loss and training accuracy metrics.\n",
    "      In the form (train_loss, train_accuracy). For example:\n",
    "\n",
    "      (0.1112, 0.8743)\n",
    "    \"\"\"\n",
    "    # Put model in train mode\n",
    "    model.train()\n",
    "\n",
    "    # Setup train loss and train accuracy values\n",
    "    train_loss, train_acc = 0, 0\n",
    "\n",
    "    # Loop through data loader data batches\n",
    "    progress_bar = tqdm(\n",
    "        enumerate(dataloader),\n",
    "        desc=f\"Training Epoch {epoch}\",\n",
    "        total=len(dataloader),\n",
    "        disable=disable_progress_bar,\n",
    "    )\n",
    "\n",
    "    for batch, (X, y) in progress_bar:\n",
    "        # Send data to target device\n",
    "        X, y = X.to(device), y.to(device)\n",
    "\n",
    "        # 1. Forward pass\n",
    "        y_pred = model(X)\n",
    "\n",
    "        # 2. Calculate  and accumulate loss\n",
    "        loss = loss_fn(y_pred, y)\n",
    "        train_loss += loss.item()\n",
    "\n",
    "        # 3. Optimizer zero grad\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # 4. Loss backward\n",
    "        loss.backward()\n",
    "\n",
    "        # 5. Optimizer step\n",
    "        optimizer.step()\n",
    "\n",
    "        # Calculate and accumulate accuracy metric across all batches\n",
    "        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)\n",
    "        train_acc += (y_pred_class == y).sum().item() / len(y_pred)\n",
    "\n",
    "        # Update progress bar\n",
    "        progress_bar.set_postfix(\n",
    "            {\n",
    "                \"train_loss\": train_loss / (batch + 1),\n",
    "                \"train_acc\": train_acc / (batch + 1),\n",
    "            }\n",
    "        )\n",
    "\n",
    "    # Adjust metrics to get average loss and accuracy per batch\n",
    "    train_loss = train_loss / len(dataloader)\n",
    "    train_acc = train_acc / len(dataloader)\n",
    "    return train_loss, train_acc\n",
    "\n",
    "\n",
    "def test_step(\n",
    "    epoch: int,\n",
    "    model: torch.nn.Module,\n",
    "    dataloader: torch.utils.data.DataLoader,\n",
    "    loss_fn: torch.nn.Module,\n",
    "    device: torch.device,\n",
    "    disable_progress_bar: bool = False,\n",
    ") -> Tuple[float, float]:\n",
    "    \"\"\"Tests a PyTorch model for a single epoch.\n",
    "\n",
    "    Turns a target PyTorch model to \"eval\" mode and then performs\n",
    "    a forward pass on a testing dataset.\n",
    "\n",
    "    Args:\n",
    "      model: A PyTorch model to be tested.\n",
    "      dataloader: A DataLoader instance for the model to be tested on.\n",
    "      loss_fn: A PyTorch loss function to calculate loss on the test data.\n",
    "      device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "      A tuple of testing loss and testing accuracy metrics.\n",
    "      In the form (test_loss, test_accuracy). For example:\n",
    "\n",
    "      (0.0223, 0.8985)\n",
    "    \"\"\"\n",
    "    # Put model in eval mode\n",
    "    model.eval()\n",
    "\n",
    "    # Setup test loss and test accuracy values\n",
    "    test_loss, test_acc = 0, 0\n",
    "\n",
    "    # Loop through data loader data batches\n",
    "    progress_bar = tqdm(\n",
    "        enumerate(dataloader),\n",
    "        desc=f\"Testing Epoch {epoch}\",\n",
    "        total=len(dataloader),\n",
    "        disable=disable_progress_bar,\n",
    "    )\n",
    "\n",
    "    # Turn on inference context manager\n",
    "    with torch.no_grad():  # no_grad() required for PyTorch 2.0, I found some errors with `torch.inference_mode()`, please let me know if this is not the case\n",
    "        # Loop through DataLoader batches\n",
    "        for batch, (X, y) in progress_bar:\n",
    "            # Send data to target device\n",
    "            X, y = X.to(device), y.to(device)\n",
    "\n",
    "            # 1. Forward pass\n",
    "            test_pred_logits = model(X)\n",
    "\n",
    "            # 2. Calculate and accumulate loss\n",
    "            loss = loss_fn(test_pred_logits, y)\n",
    "            test_loss += loss.item()\n",
    "\n",
    "            # Calculate and accumulate accuracy\n",
    "            test_pred_labels = test_pred_logits.argmax(dim=1)\n",
    "            test_acc += (test_pred_labels == y).sum().item() / len(test_pred_labels)\n",
    "\n",
    "            # Update progress bar\n",
    "            progress_bar.set_postfix(\n",
    "                {\n",
    "                    \"test_loss\": test_loss / (batch + 1),\n",
    "                    \"test_acc\": test_acc / (batch + 1),\n",
    "                }\n",
    "            )\n",
    "\n",
    "    # Adjust metrics to get average loss and accuracy per batch\n",
    "    test_loss = test_loss / len(dataloader)\n",
    "    test_acc = test_acc / len(dataloader)\n",
    "    return test_loss, test_acc\n",
    "\n",
    "\n",
    "def train(\n",
    "    model: torch.nn.Module,\n",
    "    train_dataloader: torch.utils.data.DataLoader,\n",
    "    test_dataloader: torch.utils.data.DataLoader,\n",
    "    optimizer: torch.optim.Optimizer,\n",
    "    loss_fn: torch.nn.Module,\n",
    "    epochs: int,\n",
    "    device: torch.device,\n",
    "    disable_progress_bar: bool = False,\n",
    ") -> Dict[str, List]:\n",
    "    \"\"\"Trains and tests a PyTorch model.\n",
    "\n",
    "    Passes a target PyTorch models through train_step() and test_step()\n",
    "    functions for a number of epochs, training and testing the model\n",
    "    in the same epoch loop.\n",
    "\n",
    "    Calculates, prints and stores evaluation metrics throughout.\n",
    "\n",
    "    Args:\n",
    "      model: A PyTorch model to be trained and tested.\n",
    "      train_dataloader: A DataLoader instance for the model to be trained on.\n",
    "      test_dataloader: A DataLoader instance for the model to be tested on.\n",
    "      optimizer: A PyTorch optimizer to help minimize the loss function.\n",
    "      loss_fn: A PyTorch loss function to calculate loss on both datasets.\n",
    "      epochs: An integer indicating how many epochs to train for.\n",
    "      device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "      A dictionary of training and testing loss as well as training and\n",
    "      testing accuracy metrics. Each metric has a value in a list for\n",
    "      each epoch.\n",
    "      In the form: {train_loss: [...],\n",
    "                    train_acc: [...],\n",
    "                    test_loss: [...],\n",
    "                    test_acc: [...]}\n",
    "      For example if training for epochs=2:\n",
    "                   {train_loss: [2.0616, 1.0537],\n",
    "                    train_acc: [0.3945, 0.3945],\n",
    "                    test_loss: [1.2641, 1.5706],\n",
    "                    test_acc: [0.3400, 0.2973]}\n",
    "    \"\"\"\n",
    "    # Create empty results dictionary\n",
    "    results = {\n",
    "        \"train_loss\": [],\n",
    "        \"train_acc\": [],\n",
    "        \"test_loss\": [],\n",
    "        \"test_acc\": [],\n",
    "        \"train_epoch_time\": [],\n",
    "        \"test_epoch_time\": [],\n",
    "    }\n",
    "\n",
    "    # Loop through training and testing steps for a number of epochs\n",
    "    for epoch in tqdm(range(epochs), disable=disable_progress_bar):\n",
    "        # Perform training step and time it\n",
    "        train_epoch_start_time = time.time()\n",
    "        train_loss, train_acc = train_step(\n",
    "            epoch=epoch,\n",
    "            model=model,\n",
    "            dataloader=train_dataloader,\n",
    "            loss_fn=loss_fn,\n",
    "            optimizer=optimizer,\n",
    "            device=device,\n",
    "            disable_progress_bar=disable_progress_bar,\n",
    "        )\n",
    "        train_epoch_end_time = time.time()\n",
    "        train_epoch_time = train_epoch_end_time - train_epoch_start_time\n",
    "\n",
    "        # Perform testing step and time it\n",
    "        test_epoch_start_time = time.time()\n",
    "        test_loss, test_acc = test_step(\n",
    "            epoch=epoch,\n",
    "            model=model,\n",
    "            dataloader=test_dataloader,\n",
    "            loss_fn=loss_fn,\n",
    "            device=device,\n",
    "            disable_progress_bar=disable_progress_bar,\n",
    "        )\n",
    "        test_epoch_end_time = time.time()\n",
    "        test_epoch_time = test_epoch_end_time - test_epoch_start_time\n",
    "\n",
    "        # Print out what's happening\n",
    "        print(\n",
    "            f\"Epoch: {epoch+1} | \"\n",
    "            f\"train_loss: {train_loss:.4f} | \"\n",
    "            f\"train_acc: {train_acc:.4f} | \"\n",
    "            f\"test_loss: {test_loss:.4f} | \"\n",
    "            f\"test_acc: {test_acc:.4f} | \"\n",
    "            f\"train_epoch_time: {train_epoch_time:.4f} | \"\n",
    "            f\"test_epoch_time: {test_epoch_time:.4f}\"\n",
    "        )\n",
    "\n",
    "        # Update results dictionary\n",
    "        results[\"train_loss\"].append(train_loss)\n",
    "        results[\"train_acc\"].append(train_acc)\n",
    "        results[\"test_loss\"].append(test_loss)\n",
    "        results[\"test_acc\"].append(test_acc)\n",
    "        results[\"train_epoch_time\"].append(train_epoch_time)\n",
    "        results[\"test_epoch_time\"].append(test_epoch_time)\n",
    "\n",
    "    # Return the filled results at the end of the epochs\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from tqdm.auto import tqdm\n",
    "from typing import Dict, List, Tuple\n",
    "\n",
    "\n",
    "def train_step(\n",
    "    epoch: int,\n",
    "    model: torch.nn.Module,\n",
    "    dataloader: torch.utils.data.DataLoader,\n",
    "    loss_fn: torch.nn.Module,\n",
    "    optimizer: torch.optim.Optimizer,\n",
    "    device: torch.device,\n",
    "    disable_progress_bar: bool = False,\n",
    ") -> Tuple[float, float]:\n",
    "    \"\"\"Trains a PyTorch model for a single epoch.\n",
    "\n",
    "    Turns a target PyTorch model to training mode and then\n",
    "    runs through all of the required training steps (forward\n",
    "    pass, loss calculation, optimizer step).\n",
    "\n",
    "    Args:\n",
    "      model: A PyTorch model to be trained.\n",
    "      dataloader: A DataLoader instance for the model to be trained on.\n",
    "      loss_fn: A PyTorch loss function to minimize.\n",
    "      optimizer: A PyTorch optimizer to help minimize the loss function.\n",
    "      device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "      A tuple of training loss and training accuracy metrics.\n",
    "      In the form (train_loss, train_accuracy). For example:\n",
    "\n",
    "      (0.1112, 0.8743)\n",
    "    \"\"\"\n",
    "    # Put model in train mode\n",
    "    model.train()\n",
    "\n",
    "    # Setup train loss and train accuracy values\n",
    "    train_loss, train_acc = 0, 0\n",
    "\n",
    "    # Loop through data loader data batches\n",
    "    progress_bar = tqdm(\n",
    "        enumerate(dataloader),\n",
    "        desc=f\"Training Epoch {epoch}\",\n",
    "        total=len(dataloader),\n",
    "        disable=disable_progress_bar,\n",
    "    )\n",
    "\n",
    "    for batch, (X, y) in progress_bar:\n",
    "        # Send data to target device\n",
    "        X, y = X.to(device), y.to(device)\n",
    "\n",
    "        # 1. Forward pass\n",
    "        y_pred = model(X)\n",
    "\n",
    "        # 2. Calculate  and accumulate loss\n",
    "        loss = loss_fn(y_pred, y)\n",
    "        train_loss += loss.item()\n",
    "\n",
    "        # 3. Optimizer zero grad\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # 4. Loss backward\n",
    "        loss.backward()\n",
    "\n",
    "        # 5. Optimizer step\n",
    "        optimizer.step()\n",
    "\n",
    "        # Calculate and accumulate accuracy metric across all batches\n",
    "        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)\n",
    "        train_acc += (y_pred_class == y).sum().item() / len(y_pred)\n",
    "\n",
    "        # Update progress bar\n",
    "        progress_bar.set_postfix(\n",
    "            {\n",
    "                \"train_loss\": train_loss / (batch + 1),\n",
    "                \"train_acc\": train_acc / (batch + 1),\n",
    "            }\n",
    "        )\n",
    "\n",
    "    # Adjust metrics to get average loss and accuracy per batch\n",
    "    train_loss = train_loss / len(dataloader)\n",
    "    train_acc = train_acc / len(dataloader)\n",
    "    return train_loss, train_acc\n",
    "\n",
    "\n",
    "def test_step(\n",
    "    epoch: int,\n",
    "    model: torch.nn.Module,\n",
    "    dataloader: torch.utils.data.DataLoader,\n",
    "    loss_fn: torch.nn.Module,\n",
    "    device: torch.device,\n",
    "    disable_progress_bar: bool = False,\n",
    ") -> Tuple[float, float]:\n",
    "    \"\"\"Tests a PyTorch model for a single epoch.\n",
    "\n",
    "    Turns a target PyTorch model to \"eval\" mode and then performs\n",
    "    a forward pass on a testing dataset.\n",
    "\n",
    "    Args:\n",
    "      model: A PyTorch model to be tested.\n",
    "      dataloader: A DataLoader instance for the model to be tested on.\n",
    "      loss_fn: A PyTorch loss function to calculate loss on the test data.\n",
    "      device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "      A tuple of testing loss and testing accuracy metrics.\n",
    "      In the form (test_loss, test_accuracy). For example:\n",
    "\n",
    "      (0.0223, 0.8985)\n",
    "    \"\"\"\n",
    "    # Put model in eval mode\n",
    "    model.eval()\n",
    "\n",
    "    # Setup test loss and test accuracy values\n",
    "    test_loss, test_acc = 0, 0\n",
    "\n",
    "    # Loop through data loader data batches\n",
    "    progress_bar = tqdm(\n",
    "        enumerate(dataloader),\n",
    "        desc=f\"Testing Epoch {epoch}\",\n",
    "        total=len(dataloader),\n",
    "        disable=disable_progress_bar,\n",
    "    )\n",
    "\n",
    "    # Turn on inference context manager\n",
    "    with torch.no_grad():  # no_grad() required for PyTorch 2.0, I found some errors with `torch.inference_mode()`, please let me know if this is not the case\n",
    "        # Loop through DataLoader batches\n",
    "        for batch, (X, y) in progress_bar:\n",
    "            # Send data to target device\n",
    "            X, y = X.to(device), y.to(device)\n",
    "\n",
    "            # 1. Forward pass\n",
    "            test_pred_logits = model(X)\n",
    "\n",
    "            # 2. Calculate and accumulate loss\n",
    "            loss = loss_fn(test_pred_logits, y)\n",
    "            test_loss += loss.item()\n",
    "\n",
    "            # Calculate and accumulate accuracy\n",
    "            test_pred_labels = test_pred_logits.argmax(dim=1)\n",
    "            test_acc += (test_pred_labels == y).sum().item() / len(test_pred_labels)\n",
    "\n",
    "            # Update progress bar\n",
    "            progress_bar.set_postfix(\n",
    "                {\n",
    "                    \"test_loss\": test_loss / (batch + 1),\n",
    "                    \"test_acc\": test_acc / (batch + 1),\n",
    "                }\n",
    "            )\n",
    "\n",
    "    # Adjust metrics to get average loss and accuracy per batch\n",
    "    test_loss = test_loss / len(dataloader)\n",
    "    test_acc = test_acc / len(dataloader)\n",
    "    return test_loss, test_acc\n",
    "\n",
    "\n",
    "def train(\n",
    "    model: torch.nn.Module,\n",
    "    train_dataloader: torch.utils.data.DataLoader,\n",
    "    test_dataloader: torch.utils.data.DataLoader,\n",
    "    optimizer: torch.optim.Optimizer,\n",
    "    loss_fn: torch.nn.Module,\n",
    "    epochs: int,\n",
    "    device: torch.device,\n",
    "    disable_progress_bar: bool = False,\n",
    ") -> Dict[str, List]:\n",
    "    \"\"\"Trains and tests a PyTorch model.\n",
    "\n",
    "    Passes a target PyTorch models through train_step() and test_step()\n",
    "    functions for a number of epochs, training and testing the model\n",
    "    in the same epoch loop.\n",
    "\n",
    "    Calculates, prints and stores evaluation metrics throughout.\n",
    "\n",
    "    Args:\n",
    "      model: A PyTorch model to be trained and tested.\n",
    "      train_dataloader: A DataLoader instance for the model to be trained on.\n",
    "      test_dataloader: A DataLoader instance for the model to be tested on.\n",
    "      optimizer: A PyTorch optimizer to help minimize the loss function.\n",
    "      loss_fn: A PyTorch loss function to calculate loss on both datasets.\n",
    "      epochs: An integer indicating how many epochs to train for.\n",
    "      device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "      A dictionary of training and testing loss as well as training and\n",
    "      testing accuracy metrics. Each metric has a value in a list for\n",
    "      each epoch.\n",
    "      In the form: {train_loss: [...],\n",
    "                    train_acc: [...],\n",
    "                    test_loss: [...],\n",
    "                    test_acc: [...]}\n",
    "      For example if training for epochs=2:\n",
    "                   {train_loss: [2.0616, 1.0537],\n",
    "                    train_acc: [0.3945, 0.3945],\n",
    "                    test_loss: [1.2641, 1.5706],\n",
    "                    test_acc: [0.3400, 0.2973]}\n",
    "    \"\"\"\n",
    "    # Create empty results dictionary\n",
    "    results = {\n",
    "        \"train_loss\": [],\n",
    "        \"train_acc\": [],\n",
    "        \"test_loss\": [],\n",
    "        \"test_acc\": [],\n",
    "        \"train_epoch_time\": [],\n",
    "        \"test_epoch_time\": [],\n",
    "    }\n",
    "\n",
    "    # Loop through training and testing steps for a number of epochs\n",
    "    for epoch in tqdm(range(epochs), disable=disable_progress_bar):\n",
    "        # Perform training step and time it\n",
    "        train_epoch_start_time = time.time()\n",
    "        train_loss, train_acc = train_step(\n",
    "            epoch=epoch,\n",
    "            model=model,\n",
    "            dataloader=train_dataloader,\n",
    "            loss_fn=loss_fn,\n",
    "            optimizer=optimizer,\n",
    "            device=device,\n",
    "            disable_progress_bar=disable_progress_bar,\n",
    "        )\n",
    "        train_epoch_end_time = time.time()\n",
    "        train_epoch_time = train_epoch_end_time - train_epoch_start_time\n",
    "\n",
    "        # Perform testing step and time it\n",
    "        test_epoch_start_time = time.time()\n",
    "        test_loss, test_acc = test_step(\n",
    "            epoch=epoch,\n",
    "            model=model,\n",
    "            dataloader=test_dataloader,\n",
    "            loss_fn=loss_fn,\n",
    "            device=device,\n",
    "            disable_progress_bar=disable_progress_bar,\n",
    "        )\n",
    "        test_epoch_end_time = time.time()\n",
    "        test_epoch_time = test_epoch_end_time - test_epoch_start_time\n",
    "\n",
    "        # Print out what's happening\n",
    "        print(\n",
    "            f\"Epoch: {epoch+1} | \"\n",
    "            f\"train_loss: {train_loss:.4f} | \"\n",
    "            f\"train_acc: {train_acc:.4f} | \"\n",
    "            f\"test_loss: {test_loss:.4f} | \"\n",
    "            f\"test_acc: {test_acc:.4f} | \"\n",
    "            f\"train_epoch_time: {train_epoch_time:.4f} | \"\n",
    "            f\"test_epoch_time: {test_epoch_time:.4f}\"\n",
    "        )\n",
    "\n",
    "        # Update results dictionary\n",
    "        results[\"train_loss\"].append(train_loss)\n",
    "        results[\"train_acc\"].append(train_acc)\n",
    "        results[\"test_loss\"].append(test_loss)\n",
    "        results[\"test_acc\"].append(test_acc)\n",
    "        results[\"train_epoch_time\"].append(train_epoch_time)\n",
    "        results[\"test_epoch_time\"].append(test_epoch_time)\n",
    "\n",
    "    # Return the filled results at the end of the epochs\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the number of epochs as a constant\n",
    "NUM_EPOCHS = 5\n",
    "\n",
    "# Set the learning rate as a constant (this can be changed to get better results but for now we're just focused on time)\n",
    "LEARNING_RATE = 0.003"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Epoch 0: 100%|██████████| 1563/1563 [07:25<00:00,  3.51it/s, train_loss=1.13, train_acc=0.599]\n",
      "Testing Epoch 0: 100%|██████████| 313/313 [00:24<00:00, 12.96it/s, test_loss=0.787, test_acc=0.728]\n",
      " 20%|██        | 1/5 [07:49<31:19, 469.97s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1 | train_loss: 1.1275 | train_acc: 0.5993 | test_loss: 0.7871 | test_acc: 0.7277 | train_epoch_time: 445.5800 | test_epoch_time: 24.3934\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    }
   ],
   "source": [
    "# Create model\n",
    "model, transforms = create_model()\n",
    "model.to(device)\n",
    "\n",
    "# Create loss function and optimizer\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)\n",
    "\n",
    "# Train model and track results\n",
    "single_run_no_compile_results = train(\n",
    "    model=model,\n",
    "    train_dataloader=train_dataloader,\n",
    "    test_dataloader=test_dataloader,\n",
    "    loss_fn=loss_fn,\n",
    "    optimizer=optimizer,\n",
    "    epochs=NUM_EPOCHS,\n",
    "    device=device,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create filename to save the results\n",
    "DATASET_NAME = \"CIFAR10\"\n",
    "MODEL_NAME = \"ResNet50\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "single_run_no_compile_results_df = pd.DataFrame(single_run_no_compile_results)\n",
    "# Make a directory for single_run results\n",
    "import os\n",
    "\n",
    "pytorch_2_results_dir = \"pytorch_2_results\"\n",
    "pytorch_2_single_run_results_dir = f\"{pytorch_2_results_dir}/single_run_results\"\n",
    "os.makedirs(pytorch_2_single_run_results_dir, exist_ok=True)\n",
    "\n",
    "# Create filenames for each of the dataframes\n",
    "save_name_for_non_compiled_results = (\n",
    "    f\"single_run_non_compiled_results_{DATASET_NAME}_{MODEL_NAME}.csv\"\n",
    ")\n",
    "\n",
    "# Create filepaths to save the results to\n",
    "single_run_no_compile_save_path = (\n",
    "    f\"{pytorch_2_single_run_results_dir}/{save_name_for_non_compiled_results}\"\n",
    ")\n",
    "print(\n",
    "    f\"[INFO] Saving non-compiled experiment 1 results to: {single_run_no_compile_save_path}\"\n",
    ")\n",
    "\n",
    "# Save the results\n",
    "single_run_no_compile_results_df.to_csv(single_run_no_compile_save_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "runtime",
   "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.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}