examples/notebooks/classification-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Catalyst classification 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": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install -U catalyst\n",
    "!pip install albumentations\n",
    "!pip install pretrainedmodels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Colab extras\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.\n",
    "\n",
    "\n",
    "To intergate visualization library `plotly` to colab, run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "\n",
    "def configure_plotly_browser_state():\n",
    "    display(IPython.core.display.HTML('''\n",
    "        <script src=\"/static/components/requirejs/require.js\"></script>\n",
    "        <script>\n",
    "          requirejs.config({\n",
    "            paths: {\n",
    "              base: '/static/base',\n",
    "              plotly: 'https://cdn.plot.ly/plotly-latest.min.js?noext',\n",
    "            },\n",
    "          });\n",
    "        </script>\n",
    "        '''))\n",
    "\n",
    "\n",
    "IPython.get_ipython().events.register('pre_run_cell', configure_plotly_browser_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up GPUs\n",
    "PyTorch and Catalyst versions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, catalyst\n",
    "\n",
    "torch.__version__, catalyst.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also specify GPU/CPU usage for this turorial.\n",
    "\n",
    "Available GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.utils import get_available_gpus\n",
    "\n",
    "get_available_gpus()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import List, Tuple, Callable\n",
    "\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"  # \"\" - CPU, \"0\" - 1 GPU, \"0,1\" - MultiGPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reproducibility first"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Catalyst provides a special utils for research results reproducibility. <br/>\n",
    "For example, `catalyst.utils.set_global_seed` fixes seed for all main DL frameworks (` PyTorch`, `Tensorflow`,` random` and `numpy`)\n",
    "\n",
    "In case of CuDNN you can set deterministic mode and flag to use benchmark with\n",
    "`catalyst.utils.prepare_cudnn`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 42\n",
    "from catalyst.utils import set_global_seed, prepare_cudnn\n",
    "\n",
    "set_global_seed(SEED)\n",
    "prepare_cudnn(deterministic=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will use one of four datasets:\n",
    "\n",
    "- [Ants / Bees dataset](https://www.dropbox.com/s/ffzfpbwzwdo9qp8/ants_bees_cleared_190806.tar.gz )\n",
    "- [Stanford Dogs dataset](http://vision.stanford.edu/aditya86/ImageNetDogs/) \n",
    "- [Flowers Recognition](https://www.kaggle.com/alxmamaev/flowers-recognition)\n",
    "- [Best Artworks of All Time](https://www.kaggle.com/ikarus777/best-artworks-of-all-time)\n",
    "\n",
    "If you are on MacOS and you don’t have `wget`, you can install it with:` brew install wget`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's specify the dataset by `DATASET` variable:\n",
    "- `ants_bees` (~ 400 pictures for 2 classes)  – for CPU experiments\n",
    "- `flowers` (~ 4k pictures for 5 classes)     – for fast GPU experiments \n",
    "- `dogs` (~ 20k pictures for 120 classes)\n",
    "- `artworks` (~ 8k pictures for 50 classes)   – for GPU experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET = \"artworks\" # \"ants_bees\" / \"flowers\" / \"dogs\" / \"artworks\"\n",
    "os.environ[\"DATASET\"] = DATASET"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "function gdrive_download () {\n",
    "  CONFIRM=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate \"https://docs.google.com/uc?export=download&id=$1\" -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')\n",
    "  wget --load-cookies /tmp/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$CONFIRM&id=$1\" -O $2\n",
    "  rm -rf /tmp/cookies.txt\n",
    "}\n",
    "\n",
    "mkdir -p Images\n",
    "rm -rf Images/\n",
    "\n",
    "if [[ \"$DATASET\" == \"ants_bees\" ]]; then\n",
    "    gdrive_download 1czneYKcE2sT8dAMHz3FL12hOU7m1ZkE7 ants_bees_cleared_190806.tar.gz\n",
    "    tar -xf ants_bees_cleared_190806.tar.gz &>/dev/null\n",
    "    mv ants_bees_cleared_190806 Images\n",
    "elif [[ \"$DATASET\" == \"flowers\" ]]; then\n",
    "    gdrive_download 1rvZGAkdLlbR_MEd4aDvXW11KnLaVRGFM flowers.tar.gz\n",
    "    tar -xf flowers.tar.gz &>/dev/null\n",
    "    mv flowers Images\n",
    "elif [[ \"$DATASET\" == \"dogs\" ]]; then\n",
    "    # https://www.kaggle.com/jessicali9530/stanford-dogs-dataset\n",
    "    wget http://vision.stanford.edu/aditya86/ImageNetDogs/images.tar\n",
    "    tar -xf images.tar &>/dev/null\n",
    "elif [[ \"$DATASET\" == \"artworks\" ]]; then\n",
    "    # https://www.kaggle.com/ikarus777/best-artworks-of-all-time\n",
    "    gdrive_download 1eAk36MEMjKPKL5j9VWLvNTVKk4ube9Ml artworks.tar.gz\n",
    "    tar -xf artworks.tar.gz &>/dev/null\n",
    "    mv artworks Images\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "ROOT = \"Images/\"\n",
    "ALL_IMAGES = list(Path(ROOT).glob(\"**/*.jpg\"))\n",
    "ALL_IMAGES = list(filter(lambda x: not x.name.startswith(\".\"), ALL_IMAGES))\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",
    "`create_dataset` function goes through a given directory and creates a dictionary `Dict[class_name, List[image]]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.utils.dataset import create_dataset, create_dataframe, prepare_dataset_labeling\n",
    "\n",
    "dataset = create_dataset(dirs=f\"{ROOT}/*\", extension=\"*.jpg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and `create_dataframe` function creates typical `pandas.DataFrame` for further analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = create_dataframe(dataset, columns=[\"class\", \"filepath\"])\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "finally `prepare_dataset_labeling` creates a numerical label for each unique class name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tag_to_label = prepare_dataset_labeling(df, \"class\")\n",
    "tag_to_label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's add a column with a numerical label value to the DataFrame. \n",
    "It can be easily done with `map_dataframe` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from catalyst.utils.pandas import map_dataframe\n",
    "\n",
    "df_with_labels = map_dataframe(df, tag_column=\"class\", class_column=\"label\", tag2class=tag_to_label, verbose=True)\n",
    "df_with_labels.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "additionaly let's save the `class_names` for further usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names = [name for name, id_ in sorted(tag_to_label.items(), key=lambda x: x[1])]\n",
    "class_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's divide our dataset into the `train` and` valid` parts. \n",
    "\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). \n",
    "\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.dataset import split_dataframe\n",
    "\n",
    "train_data, valid_data = split_dataframe(df_with_labels, test_size=0.2, random_state=SEED)\n",
    "train_data, valid_data = train_data.to_dict('records'), valid_data.to_dict('records')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "len(train_data), len(valid_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Augmentations"
   ]
  },
  {
   "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.dl import utils\n",
    "from catalyst.data.reader import ImageReader, ScalarReader, ReaderCompose, LambdaReader\n",
    "from functools import partial\n",
    "\n",
    "import torch\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 `datapath` folder 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",
    "        datapath=ROOT\n",
    "    ),\n",
    "    \n",
    "    # Reads a number from our dataframe 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": [
    "For augmentation of our dataset, we will use the [albumentations library](https://github.com/albu/albumentations).  <br/>\n",
    "You can view the list of available augmentations on the documentation [website](https://albumentations.readthedocs.io/en/latest/api/augmentations.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from albumentations import Compose, LongestMaxSize, PadIfNeeded\n",
    "from albumentations import ShiftScaleRotate, IAAPerspective, RandomBrightnessContrast, RandomGamma, \\\n",
    "    HueSaturationValue, ToGray, CLAHE, JpegCompression\n",
    "\n",
    "from albumentations import Normalize\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",
    "        LongestMaxSize(max_size=image_size),\n",
    "        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",
    "        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",
    "        IAAPerspective(scale=(0.02, 0.05), p=0.3),\n",
    "        # Random brightness / contrast with a 30% probability\n",
    "        RandomBrightnessContrast(\n",
    "            brightness_limit=0.2, contrast_limit=0.2, p=0.3\n",
    "        ),\n",
    "        # Random gamma changes with a 30% probability\n",
    "        RandomGamma(gamma_limit=(85, 115), p=0.3),\n",
    "        # Randomly changes the hue, saturation, and color value of the input image\n",
    "        HueSaturationValue(p=0.3),\n",
    "        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 [Normalize(), ToTensor()]\n",
    "\n",
    "def compose(_transforms):\n",
    "    # combine all augmentations into one single pipeline\n",
    "    result = Compose([item for sublist in _transforms for item in sublist])\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.augmentor import Augmentor\n",
    "from torchvision import transforms\n",
    "\n",
    "\n",
    "train_transforms = compose([pre_transforms(), hard_transforms(), post_transforms()])\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` and performs `train_transforms` on it.\n",
    "train_data_transforms = transforms.Compose([\n",
    "    Augmentor(\n",
    "        dict_key=\"features\",\n",
    "        augment_fn=lambda x: train_transforms(image=x)[\"image\"]\n",
    "    )\n",
    "])\n",
    "\n",
    "\n",
    "# Similarly for the validation part of the dataset. \n",
    "# we only perform squaring, normalization and ToTensor\n",
    "valid_data_transforms = transforms.Compose([\n",
    "    Augmentor(\n",
    "        dict_key=\"features\",\n",
    "        augment_fn=lambda x: valid_transforms(image=x)[\"image\"]\n",
    "    )\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at the augmented results. <br/>\n",
    "The cell below can be restarted."
   ]
  },
  {
   "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": [
    "## Pytorch dataloaders\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",
    "bs = 64\n",
    "num_workers = 4\n",
    "\n",
    "def get_loaders(\n",
    "    open_fn: Callable,\n",
    "    train_transforms_fn: transforms.Compose,\n",
    "    valid_transforms_fn: transforms.Compose,\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=bs,\n",
    "        num_workers=num_workers,\n",
    "        shuffle=sampler is None, # shuffle data only if Sampler is not specified (PyTorch requirement)\n",
    "        sampler=sampler\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=bs,\n",
    "        num_workers=num_workers,\n",
    "        shuffle=False, \n",
    "        sampler=None\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",
    "loaders = get_loaders(open_fn, train_data_transforms, valid_data_transforms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 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. <br/>\n",
    "But at first, let's check them out!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pretrainedmodels\n",
    "\n",
    "pretrainedmodels.model_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this tutorial purposes, `ResNet18` is good enough, but you can try other models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"resnet18\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By `pretrained_settings` we can see what the given network expects as input and what would be the expected output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrainedmodels.pretrained_settings[model_name]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model returns logits for classification into 1000 classes from ImageNet. <br/>\n",
    "Let's define a function that will replace the last fully-conected layer for our number of classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\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",
    "    dim_feats = model.last_linear.in_features\n",
    "    model.last_linear = nn.Linear(dim_feats, num_classes)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# model creation\n",
    "model = get_model(model_name, num_classes)\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.runner import SupervisedRunner\n",
    "\n",
    "runner = SupervisedRunner()\n",
    "\n",
    "# folder for all the experiment logs\n",
    "logdir = \"./logs/classification_tutorial_0\"\n",
    "NUM_EPOCHS = 10"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# as we are working on classification task\n",
    "from catalyst.dl.callbacks import AccuracyCallback, AUCCallback, F1ScoreCallback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "runner.train(\n",
    "    model=model,\n",
    "    logdir=logdir,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\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=[\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",
    "    num_epochs=NUM_EPOCHS,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training analysis and model predictions \n",
    "\n",
    "The `utils.plot_metrics` method reads tensorboard logs from the logdir and plots beautiful metrics with `plotly` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# it can take a while (colab issue)\n",
    "utils.plot_metrics(\n",
    "    logdir=logdir, \n",
    "    # specify which metrics we want to plot\n",
    "    metrics=[\"loss\", \"accuracy01\", \"auc/_mean\", \"f1_score\", \"_base/lr\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The method below will help us look at the predictions of the model for each image."
   ]
  },
  {
   "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.forward(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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 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#module-catalyst.contrib.criterion.ce):\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#module-catalyst.contrib.schedulers.base):\n",
    "- OneCycleLRWithWarmup\n",
    "\n",
    "Moreover, in [catalyst.contrib.models](https://catalyst-team.github.io/catalyst/api/contrib.html#models) 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#module-catalyst.contrib.modules.common):\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": {},
   "outputs": [],
   "source": [
    "from catalyst.contrib.criterion import FocalLossMultiClass\n",
    "from catalyst.contrib.schedulers import OneCycleLRWithWarmup\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\")\n",
    "logdir = \"./logs/classification_tutorial_1\"\n",
    "NUM_EPOCHS = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "runner.train(\n",
    "    model=model,\n",
    "    logdir=logdir,\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",
    "    num_epochs=NUM_EPOCHS,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# it can take a while (colab issue)\n",
    "utils.plot_metrics(\n",
    "    logdir=logdir, \n",
    "    metrics=[\"loss\", \"accuracy01\", \"auc/_mean\", \"f1_score\", \"_base/lr\"]\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 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": {},
   "outputs": [],
   "source": [
    "from catalyst.data.sampler import BalanceClassSampler\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()\n",
    "logdir = \"./logs/classification_tutorial_2\"\n",
    "NUM_EPOCHS = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "runner.train(\n",
    "    model=model,\n",
    "    logdir=logdir,\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",
    "    num_epochs=NUM_EPOCHS,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# it can take a while (colab issue)\n",
    "utils.plot_metrics(\n",
    "    logdir=logdir, \n",
    "    metrics=[\"loss\", \"accuracy01\", \"auc/_mean\", \"f1_score\", \"_base/lr\"]\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model inference\n",
    "\n",
    "With SupervisedRunner, you can easily predict entire loader with only one method call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = runner.predict_loader(\n",
    "    model, loaders[\"valid\"],\n",
    "    resume=f\"{logdir}/checkpoints/best.pth\", verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting object has shape = (number of elements in the loader, output shape from the model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "predictions.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus, we can obtain probabilities for our classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"logits: \", predictions[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "probabilities = torch.softmax(torch.from_numpy(predictions[0]), dim=0)\n",
    "print(\"probabilities: \", probabilities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label = probabilities.argmax().item()\n",
    "print(f\"predicted: {class_names[label]}\")\n"
   ]
  }
 ],
 "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.0"
    }
  },
 "nbformat": 4,
 "nbformat_minor": 2
}