examples/notebooks/classification-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Catalyst classification tutorial\n",
    "\n",
    "Authors: [Roman Tezikov](https://github.com/TezRomacH), [Sergey Kolesnikov](https://github.com/Scitator)\n",
    "\n",
    "[![Catalyst logo](https://raw.githubusercontent.com/catalyst-team/catalyst-pics/master/pics/catalyst_logo.png)](https://github.com/catalyst-team/catalyst)\n",
    "\n",
    "### Colab setup\n",
    "\n",
    "First of all, do not forget to change the runtime type to GPU. <br/>\n",
    "To do so click `Runtime` -> `Change runtime type` -> Select `\"Python 3\"` and `\"GPU\"` -> click `Save`. <br/>\n",
    "After that you can click `Runtime` -> `Run` all and watch the tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Requirements\n",
    "\n",
    "Download and install the latest version of catalyst and other libraries required for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this variable will be used in `runner.train` and by default we disable FP16 mode\n",
    "is_fp16_used = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Catalyst \n",
    "!pip install catalyst==20.10.1\n",
    "\n",
    "# for augmentations\n",
    "!pip install albumentations==0.4.1\n",
    "\n",
    "# for pretrained models for PyTorch\n",
    "!pip install pretrainedmodels\n",
    "\n",
    "# for TTA\n",
    "# !pip install ttach\n",
    "\n",
    "# for tensorboard\n",
    "!pip install tensorflow\n",
    "%load_ext tensorboard\n",
    "\n",
    "# if Your machine support Apex FP16, uncomment this 3 lines below\n",
    "# !git clone https://github.com/NVIDIA/apex\n",
    "# !pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./apex\n",
    "# is_fp16_used = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Callable, List, Tuple \n",
    "\n",
    "import os\n",
    "import torch\n",
    "import catalyst\n",
    "\n",
    "from catalyst.dl import utils\n",
    "\n",
    "print(f\"torch: {torch.__version__}, catalyst: {catalyst.__version__}\")\n",
    "\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"  # \"\" - CPU, \"0\" - 1 GPU, \"0,1\" - MultiGPU\n",
    "\n",
    "SEED = 42\n",
    "utils.set_global_seed(SEED)\n",
    "utils.prepare_cudnn(deterministic=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a dataset we will take [Best Artworks of All Time](https://www.kaggle.com/ikarus777/best-artworks-of-all-time) - collection of Paintings of the 50 Most Influential Artists of All Time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> If you are on MacOS and you don’t have `wget`, you can install it with: `brew install wget`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "After Catalyst installation, `download-gdrive` function become available to download objects from Google Drive.\n",
    "We use it to download datasets.\n",
    "\n",
    "usage: `download-gdrive {FILE_ID} {FILENAME}`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "download-gdrive 1eAk36MEMjKPKL5j9VWLvNTVKk4ube9Ml artworks.tar.gz &>/dev/null\n",
    "extract-archive artworks.tar.gz &>/dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "ROOT = \"artworks/\"\n",
    "ALL_IMAGES = list(Path(ROOT).glob(\"**/*.jpg\"))\n",
    "ALL_IMAGES = list(filter(lambda x: not x.name.startswith(\".\"), ALL_IMAGES))\n",
    "assert len(ALL_IMAGES) == 8446\n",
    "print(\"Number of images:\", len(ALL_IMAGES))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check out the data!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.utils import imread\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def show_examples(images: List[Tuple[str, np.ndarray]]):\n",
    "    _indexes = [(i, j) for i in range(2) for j in range(2)]\n",
    "    \n",
    "    f, ax = plt.subplots(2, 2, figsize=(16, 16))\n",
    "    for (i, j), (title, img) in zip(_indexes, images):\n",
    "        ax[i, j].imshow(img)\n",
    "        ax[i, j].set_title(title)\n",
    "    f.tight_layout()\n",
    "\n",
    "def read_random_images(paths: List[Path]) -> List[Tuple[str, np.ndarray]]:\n",
    "    data = np.random.choice(paths, size=4)\n",
    "    result = []\n",
    "    for d in data:\n",
    "        title = f\"{d.parent.name}: {d.name}\"\n",
    "        _image = imread(d)\n",
    "        result.append((title, _image))\n",
    "    \n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can restart the cell below to see more examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "images = read_random_images(ALL_IMAGES)\n",
    "show_examples(images)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset preprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With Catalyst we can easily create a dataset from the following folder structure:\n",
    "```\n",
    "dataset/\n",
    "    class_1/\n",
    "        *.ext\n",
    "        ...\n",
    "    class_2/\n",
    "        *.ext\n",
    "        ...\n",
    "    ...\n",
    "    class_N/\n",
    "        *.ext\n",
    "        ...\n",
    "```\n",
    "\n",
    "First of all `create_dataset` function goes through a given directory and creates a dictionary `Dict[class_name, List[image]]`. <br/>\n",
    "Then `create_dataframe` function creates typical `pandas.DataFrame` for further analysis. <br/>\n",
    "After than `prepare_dataset_labeling` creates a numerical label for each unique class name. <br/>\n",
    "Finally, to add a column with a numerical label value to the DataFrame, we can use `map_dataframe` function.\n",
    "\n",
    "Additionaly let's save the `class_names` for further usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.utils import (\n",
    "    create_dataset, create_dataframe, get_dataset_labeling, map_dataframe\n",
    ")\n",
    "\n",
    "dataset = create_dataset(dirs=f\"{ROOT}/*\", extension=\"*.jpg\")\n",
    "df = create_dataframe(dataset, columns=[\"class\", \"filepath\"])\n",
    "\n",
    "tag_to_label = get_dataset_labeling(df, \"class\")\n",
    "class_names = [\n",
    "    name for name, id_ in sorted(tag_to_label.items(), key=lambda x: x[1])\n",
    "]\n",
    "\n",
    "df_with_labels = map_dataframe(\n",
    "    df, \n",
    "    tag_column=\"class\", \n",
    "    class_column=\"label\", \n",
    "    tag2class=tag_to_label, \n",
    "    verbose=False\n",
    ")\n",
    "df_with_labels.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's divide our dataset into the `train` and` valid` parts. <br/>\n",
    "The parameters for the split_dataframe function are the same as [sklearn.train_test_split](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html#sklearn-model-selection-train-test-split). <br/>\n",
    "We also define `test_size` (it is optional) and `random_state` for reproducibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.utils import split_dataframe_train_test\n",
    "\n",
    "train_data, valid_data = split_dataframe_train_test(\n",
    "    df_with_labels, test_size=0.2, random_state=SEED)\n",
    "train_data, valid_data = (\n",
    "    train_data.to_dict('records'), valid_data.to_dict('records')\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to save your time during data preparation/reading/writing, Catalyst provides a special abstraction – [Reader](https://catalyst-team.github.io/catalyst/api/data.html#reader). <br/> \n",
    "Reader allows you to read various structures, for example, images, strings, numerical values and perform some functions on top of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst import utils\n",
    "from catalyst.data import ScalarReader, ReaderCompose, ImageReader\n",
    "\n",
    "num_classes = len(tag_to_label)\n",
    "\n",
    "# ReaderCompose collects different Readers into one pipeline\n",
    "open_fn = ReaderCompose([\n",
    "    \n",
    "    # Reads images from the `rootpath` folder \n",
    "    # using the key `input_key =\" filepath \"` (here should be the filename)\n",
    "    # and writes it to the output dictionary by `output_key=\"features\"` key\n",
    "    ImageReader(\n",
    "        input_key=\"filepath\",\n",
    "        output_key=\"features\",\n",
    "        rootpath=ROOT\n",
    "    ),\n",
    "    \n",
    "    # Reads a number from our dataframe \n",
    "    # by the key `input_key =\" label \"` to np.long\n",
    "    # and writes it to the output dictionary by `output_key=\"targets\"` key\n",
    "    ScalarReader(\n",
    "        input_key=\"label\",\n",
    "        output_key=\"targets\",\n",
    "        default_value=-1,\n",
    "        dtype=np.int64\n",
    "    ),\n",
    "    \n",
    "    # Same as above, but with one encoding\n",
    "    ScalarReader(\n",
    "        input_key=\"label\",\n",
    "        output_key=\"targets_one_hot\",\n",
    "        default_value=-1,\n",
    "        dtype=np.int64, \n",
    "        one_hot_classes=num_classes\n",
    "    )\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Augmentations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Albumentation logo](https://albumentations.readthedocs.io/en/latest/_static/logo.png)](https://github.com/albu/albumentations)\n",
    "\n",
    "The [albumentation](https://github.com/albu/albumentations) library works with images and masks at the same time, which is what we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import albumentations as albu\n",
    "from albumentations.pytorch import ToTensor\n",
    "\n",
    "BORDER_CONSTANT = 0\n",
    "BORDER_REFLECT = 2\n",
    "\n",
    "def pre_transforms(image_size=224):\n",
    "    # Convert the image to a square of size image_size x image_size\n",
    "    # (keeping aspect ratio)\n",
    "    result = [\n",
    "        albu.LongestMaxSize(max_size=image_size),\n",
    "        albu.PadIfNeeded(image_size, image_size, border_mode=BORDER_CONSTANT)\n",
    "    ]\n",
    "    \n",
    "    return result\n",
    "\n",
    "def hard_transforms():\n",
    "    result = [\n",
    "        # Random shifts, stretches and turns with a 50% probability\n",
    "        albu.ShiftScaleRotate( \n",
    "            shift_limit=0.1,\n",
    "            scale_limit=0.1,\n",
    "            rotate_limit=15,\n",
    "            border_mode=BORDER_REFLECT,\n",
    "            p=0.5\n",
    "        ),\n",
    "        albu.IAAPerspective(scale=(0.02, 0.05), p=0.3),\n",
    "        # Random brightness / contrast with a 30% probability\n",
    "        albu.RandomBrightnessContrast(\n",
    "            brightness_limit=0.2, contrast_limit=0.2, p=0.3\n",
    "        ),\n",
    "        # Random gamma changes with a 30% probability\n",
    "        albu.RandomGamma(gamma_limit=(85, 115), p=0.3),\n",
    "        # Randomly changes the hue, saturation, and color value of the input image\n",
    "        albu.HueSaturationValue(p=0.3),\n",
    "        albu.JpegCompression(quality_lower=80),\n",
    "    ]\n",
    "    \n",
    "    return result\n",
    "\n",
    "def post_transforms():\n",
    "    # we use ImageNet image normalization\n",
    "    # and convert it to torch.Tensor\n",
    "    return [albu.Normalize(), ToTensor()]\n",
    "\n",
    "def compose(transforms_to_compose):\n",
    "    # combine all augmentations into one single pipeline\n",
    "    result = albu.Compose([\n",
    "      item for sublist in transforms_to_compose for item in sublist\n",
    "    ])\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like Reader, there is a close abstraction for handling augmentations and key-value-based dataloaders – [Augmentor](https://catalyst-team.github.io/catalyst/api/data.html#augmentor)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.data import Augmentor\n",
    "\n",
    "\n",
    "train_transforms = compose([\n",
    "    pre_transforms(), \n",
    "    hard_transforms(), \n",
    "    post_transforms()\n",
    "])\n",
    "valid_transforms = compose([pre_transforms(), post_transforms()])\n",
    "\n",
    "show_transforms = compose([pre_transforms(), hard_transforms()])\n",
    "\n",
    "# Takes an image from the input dictionary by the key `dict_key` \n",
    "# and performs `train_transforms` on it.\n",
    "train_data_transforms = Augmentor(\n",
    "    dict_key=\"features\",\n",
    "    augment_fn=lambda x: train_transforms(image=x)[\"image\"]\n",
    ")\n",
    "\n",
    "\n",
    "# Similarly for the validation part of the dataset. \n",
    "# we only perform squaring, normalization and ToTensor\n",
    "valid_data_transforms = Augmentor(\n",
    "    dict_key=\"features\",\n",
    "    augment_fn=lambda x: valid_transforms(image=x)[\"image\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at the augmented results. <br/>\n",
    "You can restart the cell below to see more examples of augmentations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = read_random_images(ALL_IMAGES)\n",
    "\n",
    "images = [\n",
    "    (title, show_transforms(image=i)[\"image\"])\n",
    "    for (title, i) in images\n",
    "]\n",
    "show_examples(images)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loaders\n",
    "\n",
    "Using `catalyst.utils.get_loader`, you can immediately get loaders only from the dataset and data-converting functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "\n",
    "def get_loaders(\n",
    "    open_fn: Callable,\n",
    "    train_transforms_fn,\n",
    "    valid_transforms_fn,\n",
    "    batch_size: int = 64, \n",
    "    num_workers: int = 4,\n",
    "    sampler = None\n",
    ") -> collections.OrderedDict:\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        open_fn: Reader for reading data from a dataframe\n",
    "        train_transforms_fn: Augmentor for train part\n",
    "        valid_transforms_fn: Augmentor for valid part\n",
    "        batch_size: batch size\n",
    "        num_workers: How many subprocesses to use to load data,\n",
    "        sampler: An object of the torch.utils.data.Sampler class \n",
    "            for the dataset data sampling strategy specification\n",
    "    \"\"\"\n",
    "    train_loader = utils.get_loader(\n",
    "        train_data,\n",
    "        open_fn=open_fn,\n",
    "        dict_transform=train_transforms_fn,\n",
    "        batch_size=batch_size,\n",
    "        num_workers=num_workers,\n",
    "        shuffle=sampler is None, # shuffle data only if Sampler is not specified (PyTorch requirement)\n",
    "        sampler=sampler,\n",
    "        drop_last=True,\n",
    "    )\n",
    "\n",
    "    valid_loader = utils.get_loader(\n",
    "        valid_data,\n",
    "        open_fn=open_fn,\n",
    "        dict_transform=valid_transforms_fn,\n",
    "        batch_size=batch_size,\n",
    "        num_workers=num_workers,\n",
    "        shuffle=False, \n",
    "        sampler=None,\n",
    "        drop_last=True,\n",
    "    )\n",
    "\n",
    "    # Catalyst expects an ordered dictionary with train/valid/infer loaders. \n",
    "    # The number of loaders can vary.\n",
    "    # For example, it can easily handle even some complex logic like:\n",
    "    # loaders[\"train_dataset1\"] = train_loader_1\n",
    "    # loaders[\"train_dataset2\"] = train_loader_2\n",
    "    # ....\n",
    "    # loaders[\"valid_1\"] = valid_loader_1\n",
    "    # loaders[\"valid_2\"] = valid_loader_2\n",
    "    # ...\n",
    "    # loaders[\"infer_1\"] = infer_loader_1\n",
    "    # loaders[\"infer_2\"] = infer_loader_2\n",
    "    # ...\n",
    "    \n",
    "    loaders = collections.OrderedDict()\n",
    "    loaders[\"train\"] = train_loader\n",
    "    loaders[\"valid\"] = valid_loader\n",
    "\n",
    "    return loaders\n",
    "\n",
    "if is_fp16_used:\n",
    "    batch_size = 128\n",
    "else:\n",
    "    batch_size = 64\n",
    "\n",
    "print(f\"batch_size: {batch_size}\")\n",
    "\n",
    "loaders = get_loaders(\n",
    "    open_fn=open_fn, \n",
    "    train_transforms_fn=train_data_transforms,\n",
    "    valid_transforms_fn=valid_data_transforms,\n",
    "    batch_size=batch_size,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiment\n",
    "### Model\n",
    "\n",
    "Let's take the classification model from [Cadene pretrain models](https://github.com/Cadene/pretrained-models.pytorch). This repository contains a huge number of pre-trained PyTorch models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "import pretrainedmodels\n",
    "\n",
    "\n",
    "def get_model(model_name: str, num_classes: int, pretrained: str = \"imagenet\"):\n",
    "    model_fn = pretrainedmodels.__dict__[model_name]\n",
    "    model = model_fn(num_classes=1000, pretrained=pretrained)\n",
    "    \n",
    "    model.fc = nn.Sequential()\n",
    "    dim_feats = model.last_linear.in_features\n",
    "    model.last_linear = nn.Linear(dim_feats, num_classes)\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "model_name = \"resnet18\"\n",
    "model = get_model(model_name, num_classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# as we are working on basic classification problem (no multi-class/multi-label)\n",
    "# let's use standard CE loss\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0003)\n",
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(\n",
    "    optimizer, milestones=[9], gamma=0.3\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run some DL experiment, Catalyst uses a [Runner](https://catalyst-team.github.io/catalyst/api/dl.html#catalyst.dl.core.runner.Runner) abstraction. <br/>\n",
    "It contains main logic about \"how\" you run the experiment and getting predictions.\n",
    "\n",
    "For supervised learning case, there is an extention for Runner – [SupervisedRunner](https://catalyst-team.github.io/catalyst/api/dl.html#module-catalyst.dl.runner.supervised), which provides additional methods like `train`, `infer` and `predict_loader`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner\n",
    "\n",
    "num_epochs = 10\n",
    "logdir = \"./logs/classification\"\n",
    "\n",
    "device = utils.get_device()\n",
    "print(f\"device: {device}\")\n",
    "\n",
    "if is_fp16_used:\n",
    "    fp16_params = dict(opt_level=\"O1\") # params for FP16\n",
    "else:\n",
    "    fp16_params = None\n",
    "\n",
    "print(f\"FP16 params: {fp16_params}\")\n",
    "\n",
    "runner = SupervisedRunner(device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using [Callbacks](https://catalyst-team.github.io/catalyst/api/dl.html#catalyst.dl.core.callback.Callback), the basic functionality of the catalyst can be expanded.\n",
    "\n",
    "A callback is a class inherited from `catalyst.dl.core.Callback` and implements one / several / all methods:\n",
    "\n",
    "```\n",
    "on_stage_start\n",
    "    --- on_epoch_start\n",
    "    ------ on_loader_start\n",
    "    --------- on_batch_start\n",
    "    --------- on_batch_end\n",
    "    ------ on_loader_end\n",
    "    --- on_epoch_end\n",
    "on_stage_end\n",
    "\n",
    "on_exception - if the code crashes with an error, you can catch it and reserve the parameters you need\n",
    "```\n",
    "\n",
    "You can find the list of standard callbacks [here](https://catalyst-team.github.io/catalyst/api/dl.html#module-catalyst.dl.callbacks.checkpoint). \n",
    "It includes callbacks such as\n",
    "- CheckpointCallback. Saves N best models in logdir\n",
    "- TensorboardLogger. Logs all metrics to tensorboard\n",
    "- EarlyStoppingCallback. Early training stop if metrics do not improve for the `patience` of epochs\n",
    "- ConfusionMatrixCallback. Plots ConfusionMatrix per epoch in tensorboard\n",
    "\n",
    "Many different metrics for classfication\n",
    "- AccuracyCallback\n",
    "- MapKCallback\n",
    "- AUCCallback\n",
    "- F1ScoreCallback\n",
    "\n",
    "segmentation\n",
    "- DiceCallback\n",
    "- IouCallback\n",
    "\n",
    "and many other callbacks, like LRFinder and MixupCallback"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Monitoring in tensorboard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you do not have a Tensorboard opened after you have run the cell below, try running the cell again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%tensorboard --logdir {logdir}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Running train-loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# as we are working on classification task\n",
    "from catalyst.dl import AccuracyCallback, AUCCallback, F1ScoreCallback\n",
    "\n",
    "callbacks = [\n",
    "    AccuracyCallback(num_classes=num_classes),\n",
    "    AUCCallback(\n",
    "        num_classes=num_classes,\n",
    "        input_key=\"targets_one_hot\",\n",
    "        class_names=class_names\n",
    "    ),\n",
    "    F1ScoreCallback(\n",
    "        input_key=\"targets_one_hot\",\n",
    "        activation=\"Softmax\"\n",
    "    )\n",
    "]\n",
    "\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    # our dataloaders\n",
    "    loaders=loaders,\n",
    "    # We can specify the callbacks list for the experiment;\n",
    "    # For this task, we will check accuracy, AUC and F1 metrics\n",
    "    callbacks=callbacks,\n",
    "    # path to save logs\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    # save our best checkpoint by AUC metric\n",
    "    main_metric=\"auc/mean\",\n",
    "    # AUC needs to be maximized.\n",
    "    minimize_metric=False,\n",
    "    # for FP16. It uses the variable from the very first cell\n",
    "    fp16=fp16_params,\n",
    "    # prints train logs\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model inference\n",
    "\n",
    "Let's look at model predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn.functional import softmax\n",
    "\n",
    "def show_prediction(\n",
    "    model: torch.nn.Module, \n",
    "    class_names: List[str], \n",
    "    titles: List[str],\n",
    "    images: List[np.ndarray],\n",
    "    device: torch.device\n",
    ") -> None:\n",
    "    with torch.no_grad():\n",
    "        tensor_ = torch.stack([\n",
    "            valid_transforms(image=image)[\"image\"]\n",
    "            for image in images\n",
    "        ]).to(device)\n",
    "        \n",
    "        \n",
    "        logits = model(tensor_)\n",
    "        probabilities = softmax(logits, dim=1)\n",
    "        predictions = probabilities.argmax(dim=1)\n",
    "    \n",
    "    images_predicted_classes = [\n",
    "        (f\"predicted: {class_names[x]} | correct: {title}\", image)\n",
    "        for x, title, image in zip(predictions, titles, images)\n",
    "    ]\n",
    "    show_examples(images_predicted_classes)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = utils.get_device()\n",
    "titles, images = list(zip(*read_random_images(ALL_IMAGES)))\n",
    "titles = list(map(lambda x: x.rsplit(\":\")[0], titles))\n",
    "show_prediction(model, class_names=class_names, titles=titles, images=images, device=device)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally with SupervisedRunner, you can easily predict entire loader with only one method call. <br/>\n",
    "The resulting object has shape = (number of elements in the loader, output shape from the model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = np.vstack(list(map(\n",
    "    lambda x: x[\"logits\"].cpu().numpy(), \n",
    "    runner.predict_loader(loader=loaders[\"valid\"], resume=f\"{logdir}/checkpoints/best.pth\")\n",
    ")))\n",
    "print(predictions.shape)\n",
    "\n",
    "print(\"logits: \", predictions[0])\n",
    "\n",
    "probabilities = torch.softmax(torch.from_numpy(predictions[0]), dim=0)\n",
    "print(\"probabilities: \", probabilities)\n",
    "\n",
    "label = probabilities.argmax().item()\n",
    "print(f\"predicted: {class_names[label]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model tracing\n",
    "\n",
    "Catalyst allows you to use Runner to make [tracing](https://pytorch.org/docs/stable/jit.html) models.\n",
    "\n",
    "> How to do this in the Config API, we wrote in [our blog (issue \\#2)](https://github.com/catalyst-team/catalyst-info#catalyst-info-2-tracing-with-torchjit)\n",
    "\n",
    "For this purpose it is necessary to pass in a method `trace ` model and a batch on which `predict_batch ` will be executed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "batch = next(iter(loaders[\"valid\"]))\n",
    "# saves to `logdir` and returns a `ScriptModule` class\n",
    "runner.trace(model=model, batch=batch, logdir=logdir, fp16=is_fp16_used)\n",
    "\n",
    "!ls {logdir}/trace/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After this, you can easily load the model and predict anything!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from catalyst.utils import trace\n",
    "\n",
    "if is_fp16_used:\n",
    "    model = trace.load_traced_model(\n",
    "        f\"{logdir}/trace/traced-forward-opt_O1.pth\", \n",
    "        device=\"cuda\", \n",
    "        opt_level=\"O1\"\n",
    "    )\n",
    "else:\n",
    "    model = trace.load_traced_model(\n",
    "        f\"{logdir}/trace/traced-forward.pth\", \n",
    "        device=\"cpu\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model_input = batch[\"features\"].to(\"cuda\" if is_fp16_used else \"cpu\")\n",
    "model(model_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Advanced: Custom Callbacks\n",
    "\n",
    "Let's count the number of predicted labels per class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import collections\n",
    "\n",
    "from catalyst.dl import Callback, CallbackOrder, IRunner\n",
    "\n",
    "\n",
    "class CustomInferCallback(Callback):\n",
    "  def __init__(self):\n",
    "    super().__init__(CallbackOrder.Internal)\n",
    "    self.class_counts = collections.defaultdict(lambda: 0)\n",
    "\n",
    "  def on_loader_start(self, runner: IRunner):\n",
    "    self.class_counts = collections.defaultdict(lambda: 0)\n",
    "\n",
    "  def on_batch_end(self, runner: IRunner):\n",
    "    # data from the Dataloader\n",
    "    # features, targets = runner.input[\"features\"], runner.input[\"targets\"]\n",
    "    logits = runner.output[\"logits\"]\n",
    "\n",
    "    labels = logits.argmax(axis=1)\n",
    "    labels = labels.cpu().detach().numpy().tolist()\n",
    "    for label in labels:\n",
    "        self.class_counts[label] += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "from catalyst.dl.callbacks import CheckpointCallback\n",
    "\n",
    "infer_loaders = {\"infer\": loaders[\"valid\"]}\n",
    "model = get_model(model_name, num_classes)\n",
    "\n",
    "device = utils.get_device()\n",
    "if is_fp16_used:\n",
    "    fp16_params = dict(opt_level=\"O1\") # params for FP16\n",
    "else:\n",
    "    fp16_params = None\n",
    "\n",
    "runner = SupervisedRunner(device=device)\n",
    "runner.infer(\n",
    "    model=model,\n",
    "    loaders=infer_loaders,\n",
    "    callbacks=OrderedDict([\n",
    "        (\"loader\", CheckpointCallback(resume=f\"{logdir}/checkpoints/best.pth\")),\n",
    "        (\"infer\", CustomInferCallback())\n",
    "    ]),\n",
    "    fp16=fp16_params,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "\n",
    "class_counts = runner.runner._callbacks[\"infer\"].class_counts\n",
    "counts = [class_counts[x] for x in range(len(class_names))]\n",
    "\n",
    "plt.figure(figsize=(20, 9))\n",
    "x = np.arange(len(counts))\n",
    "plt.bar(x + 0.5, height=counts)\n",
    "plt.xticks(x - 0.3, class_names, rotation=65)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Advanced: Training with Focal Loss and OneCycle\n",
    "\n",
    "In the `catalyst.contrib` there are a large number of different additional criterions, models, layers etc\n",
    "\n",
    "For example,\n",
    "\n",
    "[catalyst.contrib.criterion](https://catalyst-team.github.io/catalyst/api/contrib.html):\n",
    "- HuberLoss\n",
    "- CenterLoss\n",
    "- FocalLossMultiClass\n",
    "- DiceLoss / BCEDiceLoss\n",
    "- IoULoss / BCEIoULoss\n",
    "- LovaszLossBinary / LovaszLossMultiClass / LovaszLossMultiLabel\n",
    "- WingLoss\n",
    "\n",
    "Lr scheduler in [catalyst.contrib.schedulers](https://catalyst-team.github.io/catalyst/api/contrib.html):\n",
    "- OneCycleLRWithWarmup\n",
    "\n",
    "Moreover, in [catalyst.contrib.models](https://catalyst-team.github.io/catalyst/api/contrib.html) you can find various models for segmentation:\n",
    "- Unet / ResnetUnet\n",
    "- Linknet / ResnetLinknet\n",
    "- FPNUnet / ResnetFPNUnet\n",
    "- PSPnet / ResnetPSPnet\n",
    "- MobileUnet\n",
    "\n",
    "\n",
    "Finally, several handwritten modules in [catalyst.contrib.modules](https://catalyst-team.github.io/catalyst/api/contrib.html):\n",
    "- Flatten\n",
    "- TemporalAttentionPooling\n",
    "- LamaPooling\n",
    "- NoisyLinear\n",
    "- GlobalAvgPool2d / GlobalMaxPool2d\n",
    "- GlobalConcatPool2d / GlobalAttnPool2d\n",
    "\n",
    "a bunch of others\n",
    "\n",
    "\n",
    "But for now, let's take `FocalLoss` and `OneCycleLRWithWarmup` to play around."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from catalyst.contrib.nn import FocalLossMultiClass, OneCycleLRWithWarmup\n",
    "\n",
    "logdir = \"./logs/classification_tutorial_1\"\n",
    "num_epochs = 10\n",
    "\n",
    "model = get_model(model_name, num_classes)\n",
    "\n",
    "criterion = FocalLossMultiClass()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0003)\n",
    "scheduler = OneCycleLRWithWarmup(\n",
    "    optimizer, \n",
    "    num_steps=num_epochs, \n",
    "    lr_range=(0.001, 0.0001),\n",
    "    warmup_steps=1\n",
    ")\n",
    "\n",
    "# FocalLoss expects one_hot for the input\n",
    "# in our Reader function we have already created the conversion of targets in one_hot\n",
    "# so, all we need - respecify the target key name\n",
    "runner = SupervisedRunner(input_target_key=\"targets_one_hot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        AccuracyCallback(num_classes=num_classes),\n",
    "        AUCCallback(\n",
    "            num_classes=num_classes,\n",
    "            input_key=\"targets_one_hot\",\n",
    "            class_names=class_names\n",
    "        ),\n",
    "        F1ScoreCallback(\n",
    "            input_key=\"targets_one_hot\",\n",
    "            activation=\"Softmax\"\n",
    "        )\n",
    "    ],\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    main_metric=\"auc/mean\",\n",
    "    minimize_metric=False,\n",
    "    fp16=fp16_params,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "device = utils.get_device()\n",
    "titles, images = list(zip(*read_random_images(ALL_IMAGES)))\n",
    "titles = list(map(lambda x: x.rsplit(\":\")[0], titles))\n",
    "show_prediction(model, class_names=class_names, titles=titles, images=images, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Advanced:  Balancing classes in the dataset\n",
    "\n",
    "There are several useful data-sampler implementations in the `catalyst.data.sampler`. For example,\n",
    "- `BalanceClassSampler` allows you to create stratified sampling on an unbalanced dataset. <br/> A strategy can be either downsampling, upsampling or some prespeficied number of samples per class. <br/> Very important feature for every classification problem.\n",
    "- `MiniEpochSampler` allows you to split your \"very large dataset\" and sample some small portion of it every epoch. <br/> This is useful for those cases where you need to check valid metrics and save checkpoints more often. <br/> For example, your 1M images dataset can be sampled in 100k per epoch with all necessary metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from catalyst.data import BalanceClassSampler\n",
    "\n",
    "logdir = \"./logs/classification_tutorial_2\"\n",
    "num_epochs = 10\n",
    "\n",
    "labels = [x[\"label\"] for x in train_data]\n",
    "sampler = BalanceClassSampler(labels, mode=\"upsampling\")\n",
    "\n",
    "# let's re-create our loaders with BalanceClassSampler\n",
    "loader = get_loaders(open_fn, train_data_transforms, valid_data_transforms, sampler=sampler)\n",
    "\n",
    "model = get_model(model_name, num_classes)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0003)\n",
    "scheduler = OneCycleLRWithWarmup(\n",
    "    optimizer, \n",
    "    num_steps=num_epochs, \n",
    "    lr_range=(0.001, 0.0001),\n",
    "    warmup_steps=1\n",
    ")\n",
    "\n",
    "runner = SupervisedRunner()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        AccuracyCallback(num_classes=num_classes),\n",
    "        AUCCallback(\n",
    "            num_classes=num_classes,\n",
    "            input_key=\"targets_one_hot\",\n",
    "            class_names=class_names\n",
    "        ),\n",
    "        F1ScoreCallback(\n",
    "            input_key=\"targets_one_hot\",\n",
    "            activation=\"Softmax\"\n",
    "        )\n",
    "    ],\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    main_metric=\"auc/mean\",\n",
    "    minimize_metric=False,\n",
    "    fp16=fp16_params,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "device = utils.get_device()\n",
    "titles, images = list(zip(*read_random_images(ALL_IMAGES)))\n",
    "titles = list(map(lambda x: x.rsplit(\":\")[0], titles))\n",
    "show_prediction(model, class_names=class_names, titles=titles, images=images, device=device)"
   ]
  }
 ],
 "metadata": {
  "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}