fine_tuning_segnet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "fine-tuning-segnet.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<table class=\"ee-notebook-buttons\" align=\"left\">\n",
        "    <td><a target=\"_blank\"  href=\"https://github.com/crim-ca/geoimagenet/blob/master/fine_tuning_segnet.ipynb\"><img width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /> View source on GitHub</a></td>\n",
        "    <td><a target=\"_blank\"  href=\"https://nbviewer.jupyter.org/github/crim-ca/geoimagenet/blob/master/fine_tuning_segnet.ipynb\"><img width=26px src=\"https://upload.wikimedia.org/wikipedia/commons/thumb/3/38/Jupyter_logo.svg/883px-Jupyter_logo.svg.png\" />Notebook Viewer</a></td>\n",
        "    <td><a target=\"_blank\"  href=\"https://colab.research.google.com/github/crim-ca/geoimagenet/blob/master/fine_tuning_segnet.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /> Run in Google Colab</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/crim-ca/geoimagenet/blob/master/fine_tuning_segnet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "035kgZ8EwdWn"
      },
      "source": [
        "# Fine Tuning of a SegNet Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fmTElCSfdn-W"
      },
      "source": [
        "This notebook explains how to load a geoimagenet dataset, fine-tune and package a SegNet model using the [thelper library](https://github.com/plstcharles/thelper). \n",
        "\n",
        "First, we need to install the required libraries."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hbAANhpeTdE3"
      },
      "source": [
        "### Environnement Configuration\n",
        "First, make sure that you have access to a GPU:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PCszsKlHzmcg"
      },
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "  print('and then re-execute this cell.')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8oX1hVrzt2h"
      },
      "source": [
        "In order to use a GPU with your notebook, select the Runtime > Change runtime type menu, and then set the hardware accelerator drop-down to GPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZTCbXAh_dNhe"
      },
      "source": [
        "%%bash\n",
        "pip3 --quiet install torch torchvision pillow gitpython lz4 matplotlib numpy pyyaml scikit-learn six tqdm h5py opencv-python  pretrainedmodels albumentations pyyaml\n",
        "pip3 --quiet install affine geojson shapely pyproj hdf5plugin\n",
        "pip uninstall thelper\n",
        "rm -rf ./thelper-src\n",
        "git clone  https://github.com/plstcharles/thelper.git thelper-src\n",
        "pip3 install --quiet ./thelper-src"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RWxnRTp0xMcZ"
      },
      "source": [
        "We also need to install an external repo containing new models as well as useful functions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zfnyEC4XShsT"
      },
      "source": [
        "%%bash\n",
        "pip uninstall ginmodelrepo\n",
        "rm -rf ./ginmodelrepo-src\n",
        "rm -rf ./ginmodelrepo\n",
        "git clone https://github.com/crim-ca/gin-model-repo ginmodelrepo-src\n",
        "pip3 install --quiet ./ginmodelrepo-src"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z0TssCLDxcTg"
      },
      "source": [
        "Make sure thelper is imported correctly"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dQGmKjyifhmw"
      },
      "source": [
        "import os\n",
        "import torch\n",
        "import torchvision\n",
        "import thelper\n",
        "import ginmodelrepo\n",
        "import json\n",
        "\n",
        "thelper.utils.init_logger()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wTnTevFwwVqg"
      },
      "source": [
        "We download the GeoImageNet training set:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j_qgXYu1Emzb"
      },
      "source": [
        "#ginmodelrepo.util.maybe_download_and_extract('1-tjv9gznTAM_VEaCeqz9k77fJrvVoTBw','/content/deepglobe_train.zip')\n",
        "ginmodelrepo.util.maybe_download_and_extract('12_e0dDwFDDIJTNDIBLVzJovs1YhxRhpg','/content/train/gin.zip') \n",
        "# !unzip deepglobe_train.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nXH65JznPDSP"
      },
      "source": [
        "Install the tensorboard extension:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TC70lu2vnIat"
      },
      "source": [
        "# Load the TensorBoard notebook extension\n",
        "%load_ext tensorboard"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t18ppvfjtMTM"
      },
      "source": [
        "The tensorboard below will track metrics once the training is launched, it must be launched first then hit refresh"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rjxMkOZal9eL"
      },
      "source": [
        "%tensorboard --logdir /content/gin_training/output/gin_training "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wf5VFi1HO9jQ"
      },
      "source": [
        "## Training Session Configuration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o6guok64yTkV"
      },
      "source": [
        "Name your model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SfvP5X9adOcp"
      },
      "source": [
        "session_name = 'deepglobe-unet'\n",
        "dataset_name = 'deepglobe'\n",
        "DATASET_ROOT = '/content/train'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u0dDC81QLs5h"
      },
      "source": [
        "Open the training set json file in order to prepare the data loader:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H_dfRRvLWhGa"
      },
      "source": [
        "# Opening JSON file \n",
        "with open('{}/meta.json'.format(DATASET_ROOT)) as json_file: \n",
        "  meta = json.load(json_file)\n",
        "  dataset= {'path': DATASET_ROOT, ginmodelrepo.util.DATASET_DATA_KEY : meta}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8o1isCaDIZv4"
      },
      "source": [
        "print(meta['patches'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GOY0uVxML_Fe"
      },
      "source": [
        "We check the available classes in the training set, it is important to specify the right split (train or test)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qOHUwaSMXGvP"
      },
      "source": [
        "taxo_name = 'Land cover'\n",
        "split = 'train'\n",
        "class_mapping, all_classes_with_files, all_classes_names, classes_per_taxo, datasets_per_taxo = ginmodelrepo.util.get_dataset_classes(dataset, min_occurence = 5)\n",
        "class_mapping = datasets_per_taxo[taxo_name]['class_mapping']\n",
        "print(class_mapping)\n",
        "inv_class_mapping = {str(v):k for k,v in class_mapping.items()}\n",
        "if len(all_classes_with_files) == 0:\n",
        "  raise ValueError(\"No data for split '{}'\".format(split))\n",
        "dataset[ginmodelrepo.util.DATASET_DATA_KEY][ginmodelrepo.util.DATASET_DATA_PATCH_KEY]= datasets_per_taxo[taxo_name]['data']\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bsjph7luL2Fz"
      },
      "source": [
        "Optionally, we can add the background class (everything else outside the masks) or ignore the 0 value in the masks (`dontcare = -1`)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LQYP3-s8L9GW"
      },
      "source": [
        "dontcare = -1\n",
        "if dontcare is None:\n",
        "  class_mapping = {'Background': 999, **class_mapping}\n",
        "n_classes = len(class_mapping.keys())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mm47HUDNVWB6"
      },
      "source": [
        "class_mapping = {k: str(v) for k,v in class_mapping.items()}\n",
        "print(class_mapping)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vSt4JYlAME0y"
      },
      "source": [
        "We configure the segmentation task in thelper:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KwclY5X5bJLR"
      },
      "source": [
        "task_config = {\n",
        "  \"type\": \"thelper.tasks.Segmentation\",\n",
        "  \"params\": {\n",
        "      \"class_names\": list(dict(class_mapping).keys()),\n",
        "      \"input_key\": \"data\",\n",
        "      \"dontcare\": dontcare,\n",
        "      \"label_map_key\": \"mask\"\n",
        "  }\n",
        "}\n",
        "task = ginmodelrepo.util.update_class_mapping(class_mapping, task_config)\n",
        "print(task)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UuVTl7VaMVp8"
      },
      "source": [
        "We configure the training set:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kj85pr3Jbj7c"
      },
      "source": [
        "dataset[\"dontcare\"]= dontcare\n",
        "  \n",
        "datasets_train= ginmodelrepo.util.adapt_dataset_for_model_task(task, dataset, 'train')\n",
        "datasets_valid= ginmodelrepo.util.adapt_dataset_for_model_task(task, dataset, 'test')\n",
        "datasets_config = {\n",
        "  'gin_train': datasets_train, \n",
        "  'gin_valid': datasets_valid,\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RndvWq3ZfCoS"
      },
      "source": [
        "print(datasets_config)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CjBNbEEoMdl7"
      },
      "source": [
        "The base transforms will be applied by the data loader in order to pre-process the images before training. Note that it is important to precise the `target_key` for transforms applicable to the images only (and not the masks):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SBRVfzARmZoT"
      },
      "source": [
        "base_transforms= [\n",
        "  {\n",
        "      \"operation\": \"thelper.transforms.SelectChannels\",\n",
        "      \"params\": {\"channels\": [0, 1, 2]},\n",
        "      \"target_key\": \"data\", # here we specify that only the transform is applied on the data and not the mask\n",
        "  },\n",
        "  {\n",
        "      \"operation\": \"thelper.transforms.CenterCrop\",\n",
        "      \"params\": {\"size\": 128},\n",
        "  },\n",
        "  {\n",
        "      \"operation\": \"thelper.transforms.NormalizeMinMax\",\n",
        "      \"params\": {\"min\": 0.0, \"max\": 255.0},\n",
        "      \"target_key\": \"data\",\n",
        "  },\n",
        "  {\n",
        "      \"operation\": \"thelper.transforms.NormalizeZeroMeanUnitVar\",\n",
        "      \"params\": {\n",
        "          \"mean\": [0.485, 0.456, 0.406],\n",
        "          \"std\": [0.229, 0.224, 0.225]\n",
        "      },\n",
        "      \"target_key\": \"data\",\n",
        "  },\n",
        "  {\n",
        "      \"operation\": \"thelper.transforms.Transpose\",\n",
        "      \"params\": {\n",
        "          \"axes\": [2,0,1]\n",
        "      },\n",
        "      \"target_key\": \"data\",\n",
        "  },\n",
        "]\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kd-YQZmfM-7-"
      },
      "source": [
        "Data loader creation, we need to specify the train/validation split:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tXlMGXXpdjID"
      },
      "source": [
        "loaders_config = {\n",
        "    \"base_transforms\": base_transforms,\n",
        "    \"batch_size\": 16,\n",
        "    \"workers\": 4,\n",
        "    \"pin_memory\": True,\n",
        "    \"train_split\": {\n",
        "        \"gin_train\": 1.0  \n",
        "    },\n",
        "    \"valid_split\": {\n",
        "        \"gin_valid\": 1.0 \n",
        "    },\n",
        "}\n",
        "data_config = {\"datasets\": datasets_config, \"loaders\": loaders_config}  \n",
        "save_dir = None \n",
        "task, train_loader, valid_loader, test_loader = thelper.data.create_loaders(data_config, save_dir=save_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6z92gt7zji6p"
      },
      "source": [
        "### Minibatch Visualization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GorzPqyDjoDA"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import cv2 as cv\n",
        "print(f\"train minibatch count: {len(train_loader)}\")\n",
        "samples = next(iter(train_loader))\n",
        "print(f\"images tensor shape: {samples['data'].shape}\")  # BxCxHxW\n",
        "print(f\"labels tensor shape: {samples['label'].shape}\")  # Bx1  (une étiquette par image du minibatch)\n",
        "batch_size = samples['data'].shape[0]\n",
        "display_batch_size = min(8, batch_size)\n",
        "fig = plt.figure(figsize=(24, 6))\n",
        "for r, key in enumerate(['data', 'mask']):\n",
        "  for ax_idx in range(display_batch_size):\n",
        "    ax = fig.add_subplot(2, 8, ax_idx + 1 + r*8)\n",
        "    ax.grid(False)\n",
        "    ax.set_xticks([])\n",
        "    ax.set_yticks([])\n",
        "    class_name = inv_class_mapping[task.class_names[samples['label'][ax_idx].numpy()]] \n",
        "    ax.set_title(f\"{class_name}\")\n",
        "    display = samples[key][ax_idx, ...].numpy()\n",
        "    if r == 0:\n",
        "      display = display.transpose((1, 2, 0))  # CxHxW => HxWxC (tel que demandé par matplotlib)\n",
        "      display = 128 * display + 127  # on inverse la normalisation\n",
        "    display = cv.normalize(display, dst=None, alpha=0, beta=255, norm_type=cv.NORM_MINMAX, dtype=cv.CV_8U)\n",
        "    plt.imshow(display)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-fc0JUHNMi7"
      },
      "source": [
        "### Model Configuration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PP_jpZVvKvRz"
      },
      "source": [
        "We are training a UNet-SCSE (Concurrent Spatial and Channel Squeeze & Excitation), see [this reference](https://arxiv.org/pdf/1803.02579.pdf) for more details."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aoMp3A6UgcJF"
      },
      "source": [
        "model_config = {\n",
        "    \"type\": ginmodelrepo.EncoderDecoderNet,\n",
        "    \"params\": {\"num_classes\": n_classes, \"enc_type\": 'resnet18', \"dec_type\": 'unet_scse', \"num_filters\": 16,\n",
        "            \"pretrained\": True},\n",
        "}\n",
        "model = thelper.nn.create_model({\"model\": model_config}, task)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0XhEHR5COnSj"
      },
      "source": [
        "### Trainer Configuration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kYmsGo1che2t"
      },
      "source": [
        "Gather the parameters to be optimized/updated in this run. If we are finetuning we will be updating all parameters. However, if we are doing feature extract method, we will only update the parameters that we have just  initialized, i.e. the parameters with `requires_grad = True`.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bGebQRRAiSUG"
      },
      "source": [
        "base_lr = 0.001"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f7utYgG7hdix"
      },
      "source": [
        "params_to_update = model.parameters()\n",
        "n_param = 0\n",
        "name_block = dict()\n",
        "lr_blcok = dict()\n",
        "for l, (name,param) in enumerate(model.named_parameters()):\n",
        "    n_param += 1\n",
        "for l, (name,param) in enumerate(model.named_parameters()):\n",
        "    blcok = name.split('.')[0]\n",
        "    name_block[l] = name\n",
        "    lr_blcok[l] = base_lr * 10**(2*(l/n_param-1))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FNV7BenKP26e"
      },
      "source": [
        "We turn off the model update for the encoder part, only the decoder part will be fine-tuned:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jg1o548Gjf34"
      },
      "source": [
        "params_to_update = []\n",
        "#name_block = set(name_block)\n",
        "for l, (name,param) in enumerate(model.named_parameters()):\n",
        "  if l < int(n_param/2):\n",
        "    param.requires_grad = False # we are freezing the encoder part of the UNet\n",
        "  else:\n",
        "    params_to_update.append({\n",
        "        \"params\": param,\n",
        "        \"lr\": lr_blcok[l],# you can choose to apply different lr value for each layer\n",
        "    })\n",
        "  if param.requires_grad == True:\n",
        "      print(\"\\t\",name)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BlslyZReNazS"
      },
      "source": [
        "### Training Session Creation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QDiKsUu3kjkR"
      },
      "source": [
        "optimizer_config = {\n",
        "    \"loss\": {\"type\": \"torch.nn.CrossEntropyLoss\"},\n",
        "    \"optimizer\": torch.optim.SGD(params_to_update, lr=base_lr, momentum=0.9, weight_decay=0.001),\n",
        "}\n",
        "trainer_config = {\n",
        "    \"epochs\": 20,# 20 epochs \n",
        "    \"tbx\": True, # turn on the tensorboard logs\n",
        "    \"monitor\": \"accuracy\",\n",
        "    \"optimization\": optimizer_config,\n",
        "    \"metrics\": {\n",
        "        \"accuracy\": {\"type\": \"thelper.optim.Accuracy\"},\n",
        "        \"mIoU\": {\"type\": \"thelper.optim.metrics.IntersectionOverUnion\"},\n",
        "    }\n",
        "}\n",
        "loaders = (train_loader, valid_loader, test_loader)\n",
        "trainer = thelper.train.create_trainer(session_name=\"gin_training\",\n",
        "                                       save_dir=\"./gin_training\",\n",
        "                                       config={\"trainer\": trainer_config},\n",
        "                                       model=model,\n",
        "                                       task=task,\n",
        "                                       loaders=loaders)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dMDlqQpsNgCT"
      },
      "source": [
        "### Launch the training\n",
        "Execute that cell and then go back on the Tensorboard cell on top"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-XPcbFmUk8xt"
      },
      "source": [
        "outputs = trainer.train()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PfstlyiTN5Xp"
      },
      "source": [
        "### Export the Model\n",
        "Once the training is done we can package and export the best model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2qS_GXn-N3m_"
      },
      "source": [
        "config = thelper.utils.load_checkpoint('/content/gin_training/checkpoints/ckpt.best.pth')\n",
        "model_config = {\n",
        "    \"type\": config[\"model_type\"],\n",
        "    \"params\": config[\"model_params\"],\n",
        "    'weights': config[\"model\"],\n",
        "}\n",
        "config[\"model\"] = model_config\n",
        "config[\"datasets\"] = datasets_config\n",
        "config[\"loaders\"] = loaders_config\n",
        "config[\"trainer\"] = {}\n",
        "\n",
        "thelper.cli.export_model(config, '/content/export/')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6UquM7bHjzu7"
      },
      "source": [
        "### Resuming a Session"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WLiER22UkoFT"
      },
      "source": [
        "!thelper resume --ckpt-path '/content/export/deepglobe_training/deepglobe_training.export.pth'"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}