{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HR-pt-qFMTx0"
   },
   "source": [
    "# Notebook 3: Cifar10 Classification in Pytorch\n",
    "\n",
    "In this notebook, we will train an image classifier for the CIFAR-10 dataset, that you already know from exercise 6. Today, however, we will use the PyTorch framework which makes everything much more convenient!\n",
    "We will show you how to implement the deep learning pipeline in simple PyTorch. You could also, for the first time, utilize the GPUs on colab."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IBqcIeuJMTx2"
   },
   "source": [
    "## (Optional) Mount in Google Colab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zjBDaPoXpHvt"
   },
   "source": [
    "### Set up PyTorch environment in colab\n",
    "\n",
    "For your regular environment this should already have been installed in the previous notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CzL2MBcDqAsH"
   },
   "outputs": [],
   "source": [
    "# Optional: install correct libraries in google colab. Again, downgrade numpy and then RESTART the session! Runtime > Restart Session.\n",
    "# !python -m pip install \"numpy<2.0\" torch==2.2.2+cu121 torchvision==0.17.2+cu121 -f https://download.pytorch.org/whl/torch_stable.html\n",
    "# !python -m pip install torchtext==0.17.2 torchaudio==2.2.2\n",
    "# !python -m pip install tensorboard==2.9.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UP5xfxQUMTx3"
   },
   "outputs": [],
   "source": [
    "# Use the following lines if you want to use Google Colab\n",
    "# We presume you created a folder \"i2dl\" within your main drive folder, and put the exercise there.\n",
    "# NOTE: terminate all other colab sessions that use GPU!\n",
    "# NOTE 2: Make sure the correct exercise folder (e.g exercise_07) is given.\n",
    "# OPTIONAL: Enable GPU via Runtime --> Change runtime type --> GPU\n",
    "\n",
    "\"\"\"\n",
    "from google.colab import drive\n",
    "import os\n",
    "\n",
    "gdrive_path='/content/gdrive/MyDrive/i2dl/exercise_07'\n",
    "\n",
    "# This will mount your google drive under 'MyDrive'\n",
    "drive.mount('/content/gdrive', force_remount=True)\n",
    "# In order to access the files in this notebook we have to navigate to the correct folder\n",
    "os.chdir(gdrive_path)\n",
    "# Check manually if all files are present\n",
    "print(sorted(os.listdir()))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hEDWAZ7-ZA4E"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dJCiVLV5o9QO"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True' # To prevent the kernel from dying."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dvaj6myXS7nN"
   },
   "source": [
    "### Get Device\n",
    "In this exercise, we'll use PyTorch to build an image classifier for the CIFAR-10 dataset. As you know from exercise 06, processing a large set of images is quite computation extensive. Luckily, with PyTorch we're now able to make use of our GPU to significantly speed things up!\n",
    "\n",
    "In case you don't have a GPU, you can run this notebook on Google Colab where you can access a GPU for free!\n",
    "\n",
    "Of course, you can also run this notebook on your CPU only - though this is definitely not recommended.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VWgm75NnS9hr"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Pm_rTAPnpsUo"
   },
   "source": [
    "## Setup TensorBoard\n",
    "In exercise 07 you've already learned how to use TensorBoard. Let's use it again to make the debugging of our network and training process more convenient! Throughout this notebook, feel free to add further logs or visualizations to your TensorBoard!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mI_Yf3JIMTx8"
   },
   "outputs": [],
   "source": [
    "# Delete previous instances of tensorboard\n",
    "import shutil\n",
    "tensorboard_path = os.path.abspath(\"logs\")\n",
    "if os.path.exists(tensorboard_path):\n",
    "    shutil.rmtree(tensorboard_path)\n",
    "os.makedirs(tensorboard_path, exist_ok=True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V9sz_lHyqJoj"
   },
   "source": [
    "## Define your Network\n",
    "\n",
    "Do you remember the good old times when we used to implement everything in plain numpy? Luckily, these times are over and we're using PyTorch which makes everything MUCH easier!\n",
    "\n",
    "Instead of implementing your own model, solver and dataloader, all you have to do is defining a `nn.Module`.\n",
    "\n",
    "We've prepared the class `exercise_code/MyPytorchModel` for you, that you'll now finalize to build an image classifier with PyTorch.\n",
    "\n",
    "### 0. Dataset & Dataloaders\n",
    "Check out the function `prepare_data` of the `CIFAR10DataModule` class that loads the dataset, using the class `torchvision.datasets.ImageFolder` (or the previous `MemoryImageFolder` dataset from exercise 3), which is very similar to the class `ImageFolderDataset` that you implemented earlier!\n",
    "\n",
    "Implement a **transform** to pre-process the raw data (standardize it and convert it to tensors) and assign it to the variable `my_transform`. Note: On the submission server, the normalization as in the notebook 3 on data augmentation will be performed, so please make sure to use the same normalization! For convenience, we added the precomputed normalization values for you. All normalization you are defining here are tailored to your training.\n",
    "\n",
    "In pytorch we could also include the dataset and other classes in our model, but a more reasonable way is to define it outside since it usually is used across multiple projects. If you prefer the all-in-one solution, that is great as well, but here we put it separately.\n",
    "\n",
    "If you want to improve your performance, you can also perform extensive **data augmentation** here!\n",
    "\n",
    "Also check out the `DataLoader` class that is used to create  `train_dataloader` and `val_dataloader` and that is very similar to your previous implementation of the DataLoader.\n",
    "\n",
    "### 1. Define your model\n",
    "Next, let's define your model. Think about a good network architecture. You're completely free here and you can come up with any network you like! (\\*)\n",
    "\n",
    "Have a look at the documentation of `torch.nn` at https://pytorch.org/docs/stable/nn.html to learn how to use use this module to build your network!\n",
    "\n",
    "Then implement your architecture: initialize it in `__init__()` and assign it to `self.model`. This is particularly easy using `nn.Sequential()` which you only have to pass the list of your layers.\n",
    "\n",
    "To make your model customizable and support parameter search, don't use hardcoded hyperparameters - instead, pass them as dictionary `hparams` (here, `n_hidden` is the number of neurons in the hidden layer) when initializing `MyPytorchModel`.\n",
    "\n",
    "Here's an easy example:\n",
    "\n",
    "```python\n",
    "\n",
    "    class MyPytorchModel(nn.Module):\n",
    "\n",
    "        def __init__(self, hparams):\n",
    "            super().__init__()\n",
    "            self.hparams = hparams\n",
    "           \n",
    "            self.model = nn.Sequential(\n",
    "                nn.Linear(input_size, self.hparams[\"n_hidden\"]),\n",
    "                nn.ReLU(),            \n",
    "                nn.Linear(self.hparams[\"n_hidden\"], num_classes)\n",
    "            )\n",
    "\n",
    "        def forward(self, x):\n",
    "            # Forward pass\n",
    "            out = self.model(x)\n",
    "            return out\n",
    "```\n",
    "\n",
    "or\n",
    "\n",
    "```python\n",
    "\n",
    "    class MyPytorchModel(nn.Module):\n",
    "        def __init__(self, hparams):\n",
    "            super().__init__()\n",
    "            self.hparams = hparams\n",
    "           \n",
    "            # Model\n",
    "            self.linear_1 = nn.Linear(input_size, self.hparams[\"n_hidden\"])\n",
    "            self.activation = nn.ReLU()\n",
    "            self.linear_2 = nn.Linear(self.hparams[\"n_hidden\"], num_classes)\n",
    "\n",
    "        def forward(self, x):\n",
    "            # Forward pass\n",
    "            x = self.linear_1(x)\n",
    "            x = self.activation(x)\n",
    "            x = self.linear_2(x)\n",
    "            return x\n",
    "```\n",
    "\n",
    "or\n",
    "\n",
    "```python\n",
    "\n",
    "    class MyPytorchModel(nn.Module):\n",
    "        def __init__(self, hparams):\n",
    "            super().__init__()\n",
    "            self.hparams = hparams\n",
    "           \n",
    "            # Model\n",
    "            self.linear_1 = nn.Sequential(\n",
    "                nn.Linear(input_size, self.hparams[\"n_hidden\"]),\n",
    "                nn.BatchNorm1d(self.hparams[\"n_hidden\"]),\n",
    "                nn.ReLU()\n",
    "            )\n",
    "\n",
    "            self.classifier_layer = nn.Linear(self.hparams[\"n_hidden\"], num_classes)\n",
    "\n",
    "        def forward(self, x):\n",
    "            # Forward pass\n",
    "            x = self.linear_1(x)\n",
    "            x = self.classifier_layer(x)\n",
    "            return x\n",
    "```\n",
    "\n",
    "Have a look at the forward path in `forward(self, x)`, which is so easy, that you don't need to implement it yourself. As PyTorch automatically computes the gradients, that's all we need to do! No need anymore to manually calculate derivatives for the backward paths! :)\n",
    "\n",
    "\n",
    "____\n",
    "\\* *The size of your final model must be less than 20 MB, which is approximately equivalent to 5 Mio. params. Note that this limit is quite lenient, you will probably need much less parameters!*\n",
    "\n",
    "*Also, don't use convolutional layers as they've not been covered yet in the lecture and build your network with fully connected layers (```nn.Linear()```)!*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bOYbUg8lAmgU"
   },
   "source": [
    "### 2. Training & Validation Step\n",
    "Down below we've implemented the deep learning pipeline for you. Read it carefully, and see how things are implemented in PyTorch.\n",
    "Read the comments that explain each step of the pipline.\n",
    "\n",
    "But first, let's choose our hyperparameters!\n",
    "\n",
    "It could look something like this:\n",
    "\n",
    "```python\n",
    "hparams = {\n",
    "    \"batch_size\": 64,\n",
    "    \"learning_rate\": 3e-3,\n",
    "    \"n_hidden\": 180,\n",
    "    \"input_size\": 3 * 32 * 32,\n",
    "    \"num_classes\": 10,\n",
    "    \"num_workers\": 2,\n",
    "    \"device\": device,\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_3yd6SA_MTx9"
   },
   "outputs": [],
   "source": [
    "from exercise_code.MyPytorchModel import MyPytorchModel, CIFAR10DataModule\n",
    "# make sure you have downloaded the Cifar10 dataset on root: \"../datasets/cifar10\", if not, please check exercise 03.\n",
    "hparams = {}\n",
    "\n",
    "########################################################################\n",
    "# TODO: Define your hyper parameters here!                             #\n",
    "########################################################################\n",
    "\n",
    "pass\n",
    "\n",
    "########################################################################\n",
    "#                           END OF YOUR CODE                           #\n",
    "########################################################################\n",
    "\n",
    "# Make sure you downloaded the CIFAR10 dataset already when using this cell\n",
    "# since we are showcasing the pytorch inhering ImageFolderDataset that\n",
    "# doesn't automatically download our data. Check exercise 3\n",
    "\n",
    "# If you want to switch to the memory dataset instead of image folder use\n",
    "# hparams[\"loading_method\"] = 'Memory'\n",
    "# The default is hparams[\"loading_method\"] = 'Image'\n",
    "# You will notice that it takes way longer to initialize a MemoryDataset\n",
    "# method because we have to load the data points into memory all the time.\n",
    "\n",
    "# You might get warnings below if you use too few workers. Pytorch uses\n",
    "# a more sophisticated Dataloader than the one you implemented previously.\n",
    "# In particular it uses multi processing to have multiple cores work on\n",
    "# individual data samples. You can enable more than workers (default=2)\n",
    "# via\n",
    "# hparams['num_workers'] = 8\n",
    "\n",
    "# Set up the data module including your implemented transforms\n",
    "data_module = CIFAR10DataModule(hparams)\n",
    "data_module.prepare_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oulvt67yMTx9"
   },
   "source": [
    "Some tests to check whether we'll accept your model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ur0lLT_MMTx9"
   },
   "outputs": [],
   "source": [
    "model = MyPytorchModel(hparams)\n",
    "from exercise_code.Util import printModelInfo\n",
    "_ = printModelInfo(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8v-5O-O7MTx9"
   },
   "outputs": [],
   "source": [
    "################## COLAB ##################\n",
    "# This might also work with jupyter notebooks, but will most likely not function well. Use the CMD/Terminal if possible (tensorboard --logdir=./)\n",
    "\n",
    "# %load_ext tensorboard\n",
    "# %tensorboard --logdir logs --port 6006"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_uuzXMq6zjbb"
   },
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from exercise_code.MyPytorchModel import MyPytorchModel\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "\n",
    "def create_tqdm_bar(iterable, desc):\n",
    "    return tqdm(enumerate(iterable),total=len(iterable), ncols=150, desc=desc)\n",
    "\n",
    "\n",
    "def train_model(model, train_loader, val_loader, loss_func, tb_logger, epochs=10, name=\"default\"):\n",
    "    \"\"\"\n",
    "    Train the classifier for a number of epochs.\n",
    "    \"\"\"\n",
    "    loss_cutoff = len(train_loader) // 10\n",
    "    optimizer = torch.optim.Adam(model.parameters(), hparams[\"learning_rate\"])\n",
    "\n",
    "    # The scheduler is used to change the learning rate every few \"n\" steps.\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(epochs * len(train_loader) / 5), gamma=hparams.get('gamma', 0.8))\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "\n",
    "        # Training stage, where we want to update the parameters.\n",
    "        model.train()  # Set the model to training mode\n",
    "\n",
    "        training_loss = []\n",
    "        validation_loss = []\n",
    "\n",
    "        # Create a progress bar for the training loop.\n",
    "        training_loop = create_tqdm_bar(train_loader, desc=f'Training Epoch [{epoch + 1}/{epochs}]')\n",
    "        for train_iteration, batch in training_loop:\n",
    "            optimizer.zero_grad() # Reset the gradients - VERY important! Otherwise they accumulate.\n",
    "            images, labels = batch # Get the images and labels from the batch, in the fashion we defined in the dataset and dataloader.\n",
    "            images, labels = images.to(device), labels.to(device) # Send the data to the device (GPU or CPU) - it has to be the same device as the model.\n",
    "\n",
    "            # Flatten the images to a vector. This is done because the classifier expects a vector as input.\n",
    "            # Could also be done by reshaping the images in the dataset.\n",
    "            images = images.view(images.shape[0], -1)\n",
    "\n",
    "            pred = model(images) # Stage 1: Forward().\n",
    "            loss = loss_func(pred, labels) # Compute the loss over the predictions and the ground truth.\n",
    "            loss.backward()  # Stage 2: Backward().\n",
    "            optimizer.step() # Stage 3: Update the parameters.\n",
    "            scheduler.step() # Update the learning rate.\n",
    "\n",
    "\n",
    "            training_loss.append(loss.item())\n",
    "            training_loss = training_loss[-loss_cutoff:]\n",
    "\n",
    "            # Update the progress bar.\n",
    "            training_loop.set_postfix(curr_train_loss = \"{:.8f}\".format(np.mean(training_loss)),\n",
    "                                      lr = \"{:.8f}\".format(optimizer.param_groups[0]['lr'])\n",
    "            )\n",
    "\n",
    "            # Update the tensorboard logger.\n",
    "            tb_logger.add_scalar(f'classifier_{name}/train_loss', loss.item(), epoch * len(train_loader) + train_iteration)\n",
    "\n",
    "        # Validation stage, where we don't want to update the parameters. Pay attention to the classifier.eval() line\n",
    "        # and \"with torch.no_grad()\" wrapper.\n",
    "        model.eval()\n",
    "        val_loop = create_tqdm_bar(val_loader, desc=f'Validation Epoch [{epoch + 1}/{epochs}]')\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for val_iteration, batch in val_loop:\n",
    "                images, labels = batch\n",
    "                images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "                images = images.view(images.shape[0], -1)\n",
    "                pred = model(images)\n",
    "                loss = loss_func(pred, labels)\n",
    "                validation_loss.append(loss.item())\n",
    "                # Update the progress bar.\n",
    "                val_loop.set_postfix(val_loss = \"{:.8f}\".format(np.mean(validation_loss)))\n",
    "\n",
    "                # Update the tensorboard logger.\n",
    "                tb_logger.add_scalar(f'classifier_{name}/val_loss', loss.item(), epoch * len(val_loader) + val_iteration)\n",
    "\n",
    "\n",
    "# Create a tensorboard logger.\n",
    "# NOTE: In order to see the logs, run the following command in the terminal: tensorboard --logdir=./\n",
    "# Also, in order to reset the logs, delete the logs folder MANUALLY.\n",
    "\n",
    "path = \"logs\"\n",
    "num_of_runs = len(os.listdir(path)) if os.path.exists(path) else 0\n",
    "path = os.path.join(path, f'run_{num_of_runs + 1}')\n",
    "\n",
    "tb_logger = SummaryWriter(path)\n",
    "\n",
    "# Train the classifier.\n",
    "labled_train_loader = data_module.train_dataloader()\n",
    "labled_val_loader = data_module.val_dataloader()\n",
    "\n",
    "epochs = hparams.get('epochs', 4)\n",
    "loss_func = nn.CrossEntropyLoss() # The loss function we use for classification.\n",
    "model = MyPytorchModel(hparams).to(device)\n",
    "train_model(model, labled_train_loader, labled_val_loader, loss_func, tb_logger, epochs=epochs, name=\"Default\")\n",
    "\n",
    "print()\n",
    "print(\"Finished training!\")\n",
    "print(\"How did we do? Let's check the accuracy of the defaut classifier on the training and validation sets:\")\n",
    "print(f\"Training Acc: {model.getTestAcc(labled_train_loader)[1] * 100}%\")\n",
    "print(f\"Validation Acc: {model.getTestAcc(labled_val_loader)[1] * 100}%\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zAp2OTyf4_5b"
   },
   "source": [
    "Now that everything is working, feel free to play around with different architectures. As you've seen, it's really easy to define your model or do changes there.\n",
    "\n",
    "To pass this submission, you'll need **50%** accuracy.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OmEYmRT-5S-e"
   },
   "source": [
    "# Save your model & Report Test Accuracy\n",
    "\n",
    "When you've done with your **hyperparameter tuning**, have achieved **at least 50% validation accuracy** and are happy with your final model, you can save it here.\n",
    "\n",
    "Before that, we will check again whether the number of parameters is below 5 Mi and the file size is below 20 MB.\n",
    "\n",
    "When your final model is saved, we'll lastly report the test accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "S69ETKxD5TcE"
   },
   "outputs": [],
   "source": [
    "from exercise_code.Util import test_and_save\n",
    "\n",
    "test_and_save(model, data_module.val_dataloader(), data_module.test_dataloader())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YmqQD0xiMTx_"
   },
   "source": [
    "Congrats! You've now finished your first image classifier in PyTorch! Much easier than in plain numpy, right? Time to get started with some more complex neural networks - see you at the next exercise!\n",
    "\n",
    "To create a zip file with your submission, run the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9BvMW-n9MTx_"
   },
   "outputs": [],
   "source": [
    "from exercise_code.submit import submit_exercise\n",
    "\n",
    "submit_exercise('../output/exercise07')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TQbY4sYtMTx_"
   },
   "source": [
    "# Submission Instructions\n",
    "\n",
    "Congratulations! You've just built your first image classifier with PyTorch! To complete the exercise, submit your final model to our submission portal - you probably know the procedure by now.\n",
    "\n",
    "1. Go on [our submission page](https://i2dl.vc.in.tum.de/), register for an account and login. We use your matriculation number and send an email with the login details to the mail account associated. When in doubt, login into tum online and check your mails there. You will get an ID which we need in the next step.\n",
    "2. Log into [our submission page](https://i2dl.vc.in.tum.de/) with your account details and upload the `zip` file. Once successfully uploaded, you should be able to see the submitted file selectable on the top.\n",
    "3. Click on this file and run the submission script. You will get an email with your score as well as a message if you have surpassed the threshold.\n",
    "\n",
    "# Submission Goals\n",
    "\n",
    "- Goal: Successfully implement a a fully connected NN image classifier for CIFAR-10 with PyTorch.\n",
    "\n",
    "- Passing Criteria: Similar to the last exercise, there are no unit tests that check specific components of your code. The only thing that's required to pass this optional submission, is your model to reach at least **50% accuracy** on __our__ test dataset. The submission system will show you a number between 0 and 100 which corresponds to your accuracy.\n",
    "\n",
    "- You can make **$\\infty$** submissions until the end of the semester. Remember that this exercise is an __OPTIONAL SUBMISSION__ and will __not__ be counted for the bonus."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "i2dl",
   "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.11.11"
  },
  "vscode": {
   "interpreter": {
    "hash": "54970da6898dad277dbf355945c2dee7f942d2a31ec1fc1455b6d4f552d07b83"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
