{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "99e76412-5d39-4225-a6f4-99d0f2ba702d",
   "metadata": {},
   "source": [
    "The three extensions below are optional, for more information, see\n",
    "- `watermark`:  https://github.com/rasbt/watermark\n",
    "- `pycodestyle_magic`: https://github.com/mattijn/pycodestyle_magic\n",
    "- `nb_black`: https://github.com/dnanhkhoa/nb_black"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77bd01a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext watermark\n",
    "%watermark -p torch,pytorch_lightning,torchmetrics,matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4c7e37a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext pycodestyle_magic\n",
    "%flake8_on --ignore W291,W293,E703"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eb6acf7-edaf-4bb0-aa6a-1b7f2df39285",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext nb_black"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d1a520c",
   "metadata": {},
   "source": [
    "<a href=\"https://pytorch.org\"><img src=\"https://raw.githubusercontent.com/pytorch/pytorch/master/docs/source/_static/img/pytorch-logo-dark.svg\" width=\"90\"/></a> &nbsp; &nbsp;&nbsp;&nbsp;<a href=\"https://www.pytorchlightning.ai\"><img src=\"https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/docs/source/_static/images/logo.svg\" width=\"150\"/></a>\n",
    "\n",
    "# TITLE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8323564f",
   "metadata": {},
   "source": [
    "- DESCRIPTION\n",
    "\n",
    "\n",
    "### References\n",
    "\n",
    "- ???"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec5cdbb",
   "metadata": {},
   "source": [
    "## General settings and hyperparameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9b0a861",
   "metadata": {},
   "source": [
    "- Here, we specify some general hyperparameter values and general settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb1d2d02",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 256\n",
    "NUM_EPOCHS = 10\n",
    "LEARNING_RATE = 0.005\n",
    "NUM_WORKERS = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6ec6ca5-79c1-49ba-a67c-fa76b60b7717",
   "metadata": {},
   "source": [
    "- Note that using multiple workers can sometimes cause issues with too many open files in PyTorch for small datasets. If we have problems with the data loader later, try setting `NUM_WORKERS = 0` and reload the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24b6f86d",
   "metadata": {},
   "source": [
    "## Implementing a Neural Network using PyTorch Lightning's `LightningModule`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36c91549",
   "metadata": {},
   "source": [
    "- In this section, we set up the main model architecture using the `LightningModule` from PyTorch Lightning.\n",
    "- In essence, `LightningModule` is a wrapper around a PyTorch module.\n",
    "- We start with defining our neural network model in pure PyTorch, and then we use it in the `LightningModule` to get all the extra benefits that PyTorch Lightning provides."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eacca3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# UNIQUE MODEL CODE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59d6b481",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_lightningmodule/lightningmodule_classifier_basic.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24983cbb",
   "metadata": {},
   "source": [
    "## Setting up the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6198374",
   "metadata": {},
   "source": [
    "- In this section, we are going to set up our dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "951705b9",
   "metadata": {},
   "source": [
    "### Inspecting the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17958c36",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_dataset/dataset_???_check.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee5b98a9",
   "metadata": {},
   "source": [
    "### Performance baseline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f366d74e",
   "metadata": {},
   "source": [
    "- Especially for imbalanced datasets, it's pretty helpful to compute a performance baseline.\n",
    "- In classification contexts, a useful baseline is to compute the accuracy for a scenario where the model always predicts the majority class -- we want our model to be better than that!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d8ed6d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_dataset/performance_baseline.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bdd53b7",
   "metadata": {},
   "source": [
    "## A quick visual check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "084cffe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load plot_visual-check_basic.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a8b3c3b",
   "metadata": {},
   "source": [
    "### Setting up a `DataModule`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbf59787",
   "metadata": {},
   "source": [
    "- There are three main ways we can prepare the dataset for Lightning. We can\n",
    "  1. make the dataset part of the model;\n",
    "  2. set up the data loaders as usual and feed them to the fit method of a Lightning Trainer -- the Trainer is introduced in the following subsection;\n",
    "  3. create a LightningDataModule.\n",
    "- Here, we will use approach 3, which is the most organized approach. The `LightningDataModule` consists of several self-explanatory methods, as we can see below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86d43c10",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_lightningmodule/datamodule_???_basic.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce6be803",
   "metadata": {},
   "source": [
    "- Note that the `prepare_data` method is usually used for steps that only need to be executed once, for example, downloading the dataset; the `setup` method defines the dataset loading -- if we run our code in a distributed setting, this will be called on each node / GPU. \n",
    "- Next, let's initialize the `DataModule`; we use a random seed for reproducibility (so that the data set is shuffled the same way when we re-execute this code):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78f7b3a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1) \n",
    "data_module = DataModule(data_path='./data')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59834e86",
   "metadata": {},
   "source": [
    "## Training the model using the PyTorch Lightning Trainer class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bee09340",
   "metadata": {},
   "source": [
    "- Next, we initialize our model.\n",
    "- Also, we define a call back to obtain the model with the best validation set performance after training.\n",
    "- PyTorch Lightning offers [many advanced logging services](https://pytorch-lightning.readthedocs.io/en/latest/extensions/logging.html) like Weights & Biases. However, here, we will keep things simple and use the `CSVLogger`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "970ae42c",
   "metadata": {},
   "outputs": [],
   "source": [
    "pytorch_model = PyTorchModel(\n",
    "    ???\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e3238a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_lightningmodule/logger_csv_acc_basic.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5861ddcb",
   "metadata": {},
   "source": [
    "- Now it's time to train our model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccf14578",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_lightningmodule/trainer_nb_basic.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c1ea45d",
   "metadata": {},
   "source": [
    "## Evaluating the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9358c717",
   "metadata": {},
   "source": [
    "- After training, let's plot our training ACC and validation ACC using pandas, which, in turn, uses matplotlib for plotting (PS: you may want to check out [more advanced logger](https://pytorch-lightning.readthedocs.io/en/latest/extensions/logging.html) later on, which take care of it for us):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "842ed5b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_lightningmodule/logger_csv_plot_basic.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e426862",
   "metadata": {},
   "source": [
    "- The `trainer` automatically saves the model with the best validation accuracy automatically for us, we which we can load from the checkpoint via the `ckpt_path='best'` argument; below we use the `trainer` instance to evaluate the best model on the test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27ddf180",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.test(model=lightning_model, datamodule=data_module, ckpt_path='best')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b88f525",
   "metadata": {},
   "source": [
    "## Predicting labels of new data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be642c73",
   "metadata": {},
   "source": [
    "- We can use the `trainer.predict` method either on a new `DataLoader` (`trainer.predict(dataloaders=...)`) or `DataModule` (`trainer.predict(datamodule=...)`) to apply the model to new data.\n",
    "- Alternatively, we can also manually load the best model from a checkpoint as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5367ee40",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = trainer.checkpoint_callback.best_model_path\n",
    "print(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3dad29f",
   "metadata": {},
   "outputs": [],
   "source": [
    "lightning_model = LightningModel.load_from_checkpoint(path, model=pytorch_model)\n",
    "lightning_model.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92211016",
   "metadata": {},
   "source": [
    "- For simplicity, we reused our existing `pytorch_model` above. However, we could also reinitialize the `pytorch_model`, and the `.load_from_checkpoint` method would load the corresponding model weights for us from the checkpoint file.\n",
    "- Now, below is an example applying the model manually. Here, pretend that the `test_dataloader` is a new data loader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "221df0f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_lightningmodule/datamodule_testloader.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e4f69a",
   "metadata": {},
   "source": [
    "- As an internal check, if the model was loaded correctly, the test accuracy below should be identical to the test accuracy we saw earlier in the previous section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b265a5b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_acc = acc.compute()\n",
    "print(f'Test accuracy: {test_acc:.4f} ({test_acc*100:.2f}%)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68c3e00d",
   "metadata": {},
   "source": [
    "## Inspecting Failure Cases"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cc44b41",
   "metadata": {},
   "source": [
    "- In practice, it is often informative to look at failure cases like wrong predictions for particular training instances as it can give us some insights into the model behavior and dataset.\n",
    "- Inspecting failure cases can sometimes reveal interesting patterns and even highlight dataset and labeling issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12b82d48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In the case of ???, the class label mapping\n",
    "# ???\n",
    "class_dict = {???}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e298ac1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_lightningmodule/plot_failurecases_basic.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f30b8215-caef-44fd-9657-7269880a9bd0",
   "metadata": {},
   "source": [
    "- In addition to inspecting failure cases visually, it is also informative to look at which classes the model confuses the most via a confusion matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbb4ab75",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../code_lightningmodule/plot_confusion-matrix_basic.py"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}