{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PnRi7rdVuM4g"
   },
   "source": [
    "# Prolegomena\n",
    "\n",
    "Depending on whether you are running locally or on google colab uncomment the appropriate PATH cell.\n",
    "By default the local cell is uncommented."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F6QqL3ac1-nj"
   },
   "source": [
    " ## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 2673,
     "status": "ok",
     "timestamp": 1632313052752,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "cNdnu9WaIPKn"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "torch.manual_seed(42)\n",
    "\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torch.optim import Adam\n",
    "from torch.nn import CrossEntropyLoss\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "import gc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PATH (local run)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH_TO_ROOT = \".\"\n",
    "PATH_TO_DATASETS = os.path.join(\n",
    "    PATH_TO_ROOT, \"datasets\"\n",
    ")\n",
    "\n",
    "from modules.utils import Trainer\n",
    "from modules.models import DFC_LeNet_5, Deep_RetiNet\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V8C6xMfy2EHy"
   },
   "source": [
    "## Mount drive, PATH (on google colab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 19463,
     "status": "ok",
     "timestamp": 1632313072197,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "tJNgomig9wuQ",
    "outputId": "581dcc34-c581-472e-8aa1-e322a72b1161"
   },
   "outputs": [],
   "source": [
    "# from google.colab import drive\n",
    "# import os, sys\n",
    "\n",
    "# drive.mount('./mnt')\n",
    "# !ls mnt/MyDrive/Code/ -l\n",
    "\n",
    "# PATH_TO_ROOT = \"mnt/MyDrive/Code/RetiNet\"\n",
    "# PATH_TO_DATASETS = os.path.join(\n",
    "#     PATH_TO_ROOT, \"datasets\"\n",
    "# )\n",
    "\n",
    "# sys.path.append(PATH_TO_ROOT)\n",
    "# from modules.utils import Trainer\n",
    "# from modules.models import DFC_LeNet_5, Deep_RetiNet\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 13,
     "status": "ok",
     "timestamp": 1632313072197,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "zzY0KC0e9yW5",
    "outputId": "b88dd0d9-9093-4aa1-adec-79450c0c2959"
   },
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f'Using: {device}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J4XNYGrY63S0"
   },
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input panel\n",
    "\n",
    "Training parameters input section. When a limited set of options is available, a list of the possibilities is specified as a comment above the parameter definition line.\n",
    "When training a new model you can start directly from here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "executionInfo": {
     "elapsed": 10,
     "status": "ok",
     "timestamp": 1632321780784,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "FKBVq-u_SbHk"
   },
   "outputs": [],
   "source": [
    "# MNIST - FashionMNIST - SVHN\n",
    "dataset_name = \"MNIST\" #@param         \n",
    "\n",
    "# LeNet - Deep_RetiNet\n",
    "model_name = \"Deep_RetiNet\" #@param \n",
    "\n",
    "retinic_kernel_size =  7#@param\n",
    "rks = retinic_kernel_size\n",
    "\n",
    "depth =  3#@param\n",
    "\n",
    "\n",
    "optimizer = Adam \n",
    "batch_size = 128 \n",
    "start_lr = 1e-3 \n",
    "epochs = 20\n",
    "\n",
    "loss_fn = CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Datasets loading, dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 15,
     "status": "ok",
     "timestamp": 1632321781239,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "sWnGT5FA2uiZ"
   },
   "outputs": [],
   "source": [
    "IL = { \n",
    "    \"MNIST\" : datasets.MNIST,\n",
    "    \"FashionMNIST\" : datasets.FashionMNIST,\n",
    "    \"SVHN\" : datasets.SVHN,\n",
    "}\n",
    "\n",
    "\n",
    "if dataset_name == \"SVHN\":\n",
    "    # pytorch does not use the same standards for all\n",
    "    # datasets for some reasons I dont know why\n",
    "    path = os.path.join(PATH_TO_DATASETS, \"SVHN\")\n",
    "\n",
    "    trainset = IL[dataset_name](\n",
    "        path,\n",
    "        download = True,\n",
    "        split = \"train\",\n",
    "        transform = transforms.ToTensor()\n",
    "    )\n",
    "\n",
    "    testset = IL[dataset_name](\n",
    "        path,\n",
    "        download = True,\n",
    "        split = \"test\",\n",
    "        transform = transforms.ToTensor()\n",
    "    )\n",
    "\n",
    "else:\n",
    "    transform = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Pad(2)\n",
    "    ])\n",
    "\n",
    "    trainset = IL[dataset_name](\n",
    "        PATH_TO_DATASETS,\n",
    "        download = True,\n",
    "        train = True,\n",
    "        transform = transform\n",
    "    )\n",
    "\n",
    "    testset = IL[dataset_name](\n",
    "        PATH_TO_DATASETS,\n",
    "        download = True,\n",
    "        train = False,\n",
    "        transform = transform\n",
    "    )\n",
    " \n",
    "in_channels = testset[0][0].shape[0]\n",
    "train_size = trainset.__len__()\n",
    "test_size = testset.__len__()\n",
    "\n",
    "\n",
    "trainloader = DataLoader(trainset, \n",
    "                          batch_size = batch_size, \n",
    "                          shuffle = True, \n",
    "                          num_workers = 1)\n",
    "\n",
    "testloader = DataLoader(testset, \n",
    "                        batch_size = 5000, \n",
    "                        shuffle = False, \n",
    "                        num_workers = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model and optimizer istantiation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 14,
     "status": "ok",
     "timestamp": 1632321781239,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "cAi54tRA4jO_",
    "outputId": "9725f53f-12d8-49c1-f0cc-f4732332f7b2"
   },
   "outputs": [],
   "source": [
    "if model_name == \"LeNet\":\n",
    "  model = DFC_LeNet_5(in_channels).to(device)\n",
    "  model_save_name = \"DFC_LeNet_5\"\n",
    "\n",
    "elif model_name == \"Deep_RetiNet\":\n",
    "  model = Deep_RetiNet(depth, rks, in_channels).to(device)\n",
    "  model_save_name = f\"Deep_RetiNet_d{depth}_rks{rks}\"\n",
    "\n",
    "optimizer = optimizer(model.parameters(), lr=start_lr)\n",
    "\n",
    "print(model)\n",
    "\n",
    "trainer = Trainer(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B8sUfVKwKURe"
   },
   "source": [
    "## Grid search: uncomment\n",
    "In the case you need to do some grid search uncomment this cell and run it "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 6,
     "status": "ok",
     "timestamp": 1632321781240,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "5sG4cI6lKKgR"
   },
   "outputs": [],
   "source": [
    "# lrs = [1e-3, 2e-3, 3e-3]\n",
    "# batch_sizes = [128]\n",
    "# betass = [(0.7, 0.999), (0.8, 0.999), (0.9, 0.999)]\n",
    "# epss = [1e-7, 1e-8, 1e-9]\n",
    "# gc.collect()\n",
    "\n",
    "# for lr in lrs:\n",
    "#   for batch_size in batch_sizes:\n",
    "#     print(f\"Learning rate: {lr} \\t batch size: {batch_size}\")\n",
    "\n",
    "#     gs_model = LeNet_5().to(device)\n",
    "#     gs_trainer = Trainer(gs_model)\n",
    "\n",
    "#     gs_trainloader = DataLoader(trainset, \n",
    "#                             batch_size = batch_size, \n",
    "#                             shuffle = True, \n",
    "#                             num_workers = 2)\n",
    "\n",
    "#     gs_optimizer = Adam(gs_model.parameters(),\n",
    "#                        lr = lr)\n",
    "\n",
    "#     gs_trainer.train(trainloader = gs_trainloader,\n",
    "#                      validloader = testloader,\n",
    "#                      optimizer = gs_optimizer,\n",
    "#                      epochs = 1,\n",
    "#                      loss_fn = loss_fn,\n",
    "#                      retitrain = False,\n",
    "#                      plotting = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RWUAsHx1lg7b"
   },
   "source": [
    "## Training routine\n",
    "3+1+1-scheduled trainig routine\n",
    "we do somewhat of a learning rate scheduling and we avoid dropping the colab gpu memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 436
    },
    "executionInfo": {
     "elapsed": 97876,
     "status": "ok",
     "timestamp": 1632321880164,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "VKkia_uFXkTQ",
    "outputId": "2ef44cdb-626e-4453-c236-8b14e8c04880"
   },
   "outputs": [],
   "source": [
    "gc.collect()\n",
    "\n",
    "trainer.train(trainloader = trainloader,\n",
    "              validloader = testloader,\n",
    "              optimizer = optimizer,\n",
    "              loss_fn = loss_fn,\n",
    "              epochs = 5,\n",
    "              plotting = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 436
    },
    "executionInfo": {
     "elapsed": 97999,
     "status": "ok",
     "timestamp": 1632321979066,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "g4WdlmmtbCxr",
    "outputId": "c90461ab-cb2b-4cc0-c60b-cb75ced89428"
   },
   "outputs": [],
   "source": [
    "gc.collect()\n",
    "\n",
    "trainer.train(trainloader = trainloader,\n",
    "              validloader = testloader,\n",
    "              optimizer = optimizer,\n",
    "              loss_fn = loss_fn,\n",
    "              epochs = 5,\n",
    "              plotting = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 436
    },
    "executionInfo": {
     "elapsed": 98588,
     "status": "ok",
     "timestamp": 1632322077645,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "OZ3xSYHqYvFL",
    "outputId": "d971391d-13c6-48d0-dc82-c007e65ef9d1"
   },
   "outputs": [],
   "source": [
    "gc.collect()\n",
    "\n",
    "trainer.train(trainloader = trainloader,\n",
    "              validloader = testloader,\n",
    "              optimizer = optimizer,\n",
    "              loss_fn = loss_fn,\n",
    "              epochs = 5,\n",
    "              plotting = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 436
    },
    "executionInfo": {
     "elapsed": 98463,
     "status": "ok",
     "timestamp": 1632322176098,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "cPNv4DTRYtne",
    "outputId": "a9211bac-76f7-407a-93d3-031b4167b1fc"
   },
   "outputs": [],
   "source": [
    "gc.collect()\n",
    "\n",
    "for g in optimizer.param_groups:\n",
    "  g[\"lr\"] = 1e-4\n",
    "\n",
    "trainer.train(trainloader = trainloader,\n",
    "              validloader = testloader,\n",
    "              optimizer = optimizer,\n",
    "              loss_fn = loss_fn,\n",
    "              epochs = 5,\n",
    "              plotting = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 402
    },
    "executionInfo": {
     "elapsed": 58280,
     "status": "ok",
     "timestamp": 1632322235205,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "ajuGT4mAY1ce",
    "outputId": "650c3c54-d32a-4934-8517-f61a72a64df2"
   },
   "outputs": [],
   "source": [
    "gc.collect()\n",
    "\n",
    "for g in optimizer.param_groups:\n",
    "  g[\"lr\"] = 1e-5\n",
    "\n",
    "trainer.train(trainloader = trainloader,\n",
    "              validloader = testloader,\n",
    "              optimizer = optimizer,\n",
    "              loss_fn = loss_fn,\n",
    "              epochs = 3,\n",
    "              plotting = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the state dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 7,
     "status": "ok",
     "timestamp": 1632322235206,
     "user": {
      "displayName": "Rita Fioresi",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "13337569160311637673"
     },
     "user_tz": -120
    },
    "id": "bSEl82h7eFhA"
   },
   "outputs": [],
   "source": [
    "PATH_TO_SAVE = PATH_TO_ROOT + f\"/trained_models/{dataset_name}\"\n",
    "torch.save(model.state_dict(),\n",
    "            f\"{PATH_TO_SAVE}/trained_{model_save_name}_state_dict.pt\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "training",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
