training-efficientdet-parallel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Really good training pipeline for pytorch EfficientDet \n",
    "\n",
    "Hi everyone!\n",
    "\n",
    "My name is Alex Shonenkov, I am DL/NLP/CV/TS research engineer. Especially I am in Love with NLP & DL.\n",
    "\n",
    "Recently I have created kernel for this competition about Weighted Boxes Fusion:\n",
    "- [WBF approach for ensemble](https://www.kaggle.com/shonenkov/wbf-approach-for-ensemble)\n",
    "\n",
    "\n",
    "I hope it is useful for you, my friends! If you didn't read this kernel, don't forget to do it! :)\n",
    "\n",
    "\n",
    "Today I would like to share really good training pipeline for this competition using SOTA [EfficientDet: Scalable and Efficient Object Detection](https://arxiv.org/pdf/1911.09070.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main Idea\n",
    " \n",
    "I read [all public kernels about EfficientDet in kaggle community](https://www.kaggle.com/search?q=efficientdet+in%3Anotebooks) and understand that kaggle don't have really good working public kernels with good score. Why? You can see below picture about COCO AP for different architectures, I think everyone should be able to use such a strong tools EfficientDet for own research, lets do it!\n",
    "\n",
    "<img src='https://miro.medium.com/max/2400/0*ApAKUWtseHcvRV2U.png' width=400>   \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies and imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"timm-efficientdet-pytorch\")\n",
    "sys.path.insert(0, \"omegaconf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "_kg_hide-input": true,
    "_kg_hide-output": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "from datetime import datetime\n",
    "import time\n",
    "import random\n",
    "import cv2\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import albumentations as A\n",
    "import matplotlib.pyplot as plt\n",
    "from albumentations.pytorch.transforms import ToTensorV2\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from torch.utils.data import Dataset,DataLoader\n",
    "from torch.utils.data.sampler import SequentialSampler, RandomSampler\n",
    "from glob import glob\n",
    "\n",
    "SEED = 42\n",
    "\n",
    "def seed_everything(seed):\n",
    "    random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = True\n",
    "\n",
    "seed_everything(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-17-e2563b86b642>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mmarking\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread_csv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'/home/hy/dataset/gwd/train.csv'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mbboxs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmarking\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'bbox'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mapply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfromstring\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msep\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m','\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcolumn\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'x'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'y'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'w'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'h'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0mmarking\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mcolumn\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbboxs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/badeda/lib/python3.7/site-packages/pandas/io/parsers.py\u001b[0m in \u001b[0;36mparser_f\u001b[0;34m(filepath_or_buffer, sep, delimiter, header, names, index_col, usecols, squeeze, prefix, mangle_dupe_cols, dtype, engine, converters, true_values, false_values, skipinitialspace, skiprows, skipfooter, nrows, na_values, keep_default_na, na_filter, verbose, skip_blank_lines, parse_dates, infer_datetime_format, keep_date_col, date_parser, dayfirst, cache_dates, iterator, chunksize, compression, thousands, decimal, lineterminator, quotechar, quoting, doublequote, escapechar, comment, encoding, dialect, error_bad_lines, warn_bad_lines, delim_whitespace, low_memory, memory_map, float_precision)\u001b[0m\n\u001b[1;32m    683\u001b[0m         )\n\u001b[1;32m    684\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 685\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0m_read\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilepath_or_buffer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    686\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    687\u001b[0m     \u001b[0mparser_f\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__name__\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/badeda/lib/python3.7/site-packages/pandas/io/parsers.py\u001b[0m in \u001b[0;36m_read\u001b[0;34m(filepath_or_buffer, kwds)\u001b[0m\n\u001b[1;32m    461\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    462\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 463\u001b[0;31m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mparser\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnrows\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    464\u001b[0m     \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    465\u001b[0m         \u001b[0mparser\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/badeda/lib/python3.7/site-packages/pandas/io/parsers.py\u001b[0m in \u001b[0;36mread\u001b[0;34m(self, nrows)\u001b[0m\n\u001b[1;32m   1152\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnrows\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1153\u001b[0m         \u001b[0mnrows\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_validate_integer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"nrows\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnrows\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1154\u001b[0;31m         \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_engine\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnrows\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1155\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1156\u001b[0m         \u001b[0;31m# May alter columns / col_dict\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/badeda/lib/python3.7/site-packages/pandas/io/parsers.py\u001b[0m in \u001b[0;36mread\u001b[0;34m(self, nrows)\u001b[0m\n\u001b[1;32m   2057\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnrows\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2058\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2059\u001b[0;31m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_reader\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnrows\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2060\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mStopIteration\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2061\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_first_chunk\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/parsers.pyx\u001b[0m in \u001b[0;36mpandas._libs.parsers.TextReader.read\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/parsers.pyx\u001b[0m in \u001b[0;36mpandas._libs.parsers.TextReader._read_low_memory\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/parsers.pyx\u001b[0m in \u001b[0;36mpandas._libs.parsers.TextReader._read_rows\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/parsers.pyx\u001b[0m in \u001b[0;36mpandas._libs.parsers.TextReader._convert_column_data\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/parsers.pyx\u001b[0m in \u001b[0;36mpandas._libs.parsers.TextReader._convert_tokens\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mpandas/_libs/parsers.pyx\u001b[0m in \u001b[0;36mpandas._libs.parsers.TextReader._convert_with_dtype\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/badeda/lib/python3.7/site-packages/pandas/core/dtypes/common.py\u001b[0m in \u001b[0;36mis_categorical_dtype\u001b[0;34m(arr_or_dtype)\u001b[0m\n\u001b[1;32m    676\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    677\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 678\u001b[0;31m \u001b[0;32mdef\u001b[0m \u001b[0mis_categorical_dtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marr_or_dtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    679\u001b[0m     \"\"\"\n\u001b[1;32m    680\u001b[0m     \u001b[0mCheck\u001b[0m \u001b[0mwhether\u001b[0m \u001b[0man\u001b[0m \u001b[0marray\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mlike\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mdtype\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mof\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mCategorical\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "marking = pd.read_csv('/home/hy/dataset/gwd/train.csv')\n",
    "\n",
    "bboxs = np.stack(marking['bbox'].apply(lambda x: np.fromstring(x[1:-1], sep=',')))\n",
    "for i, column in enumerate(['x', 'y', 'w', 'h']):\n",
    "    marking[column] = bboxs[:,i]\n",
    "marking.drop(columns=['bbox'], inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "About data splitting you can read [here](https://www.kaggle.com/shonenkov/wbf-approach-for-ensemble):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_kg_hide-output": true
   },
   "outputs": [],
   "source": [
    "skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\n",
    "\n",
    "df_folds = marking[['image_id']].copy()\n",
    "df_folds.loc[:, 'bbox_count'] = 1\n",
    "df_folds = df_folds.groupby('image_id').count()\n",
    "df_folds.loc[:, 'source'] = marking[['image_id', 'source']].groupby('image_id').min()['source']\n",
    "df_folds.loc[:, 'stratify_group'] = np.char.add(\n",
    "    df_folds['source'].values.astype(str),\n",
    "    df_folds['bbox_count'].apply(lambda x: f'_{x // 15}').values.astype(str)\n",
    ")\n",
    "df_folds.loc[:, 'fold'] = 0\n",
    "\n",
    "for fold_number, (train_index, val_index) in enumerate(skf.split(X=df_folds.index, y=df_folds['stratify_group'])):\n",
    "    df_folds.loc[df_folds.iloc[val_index].index, 'fold'] = fold_number"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Albumentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_train_transforms():\n",
    "    return A.Compose(\n",
    "        [\n",
    "            A.RandomSizedCrop(min_max_height=(800, 800), height=1024, width=1024, p=0.5),\n",
    "            A.OneOf([\n",
    "                A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit= 0.2, \n",
    "                                     val_shift_limit=0.2, p=0.9),\n",
    "                A.RandomBrightnessContrast(brightness_limit=0.2, \n",
    "                                           contrast_limit=0.2, p=0.9),\n",
    "            ],p=0.9),\n",
    "            A.ToGray(p=0.01),\n",
    "            A.HorizontalFlip(p=0.5),\n",
    "            A.VerticalFlip(p=0.5),\n",
    "            A.Resize(height=512, width=512, p=1),\n",
    "            A.Cutout(num_holes=8, max_h_size=64, max_w_size=64, fill_value=0, p=0.5),\n",
    "            ToTensorV2(p=1.0),\n",
    "        ], \n",
    "        p=1.0, \n",
    "        bbox_params=A.BboxParams(\n",
    "            format='pascal_voc',\n",
    "            min_area=0, \n",
    "            min_visibility=0,\n",
    "            label_fields=['labels']\n",
    "        )\n",
    "    )\n",
    "\n",
    "def get_valid_transforms():\n",
    "    return A.Compose(\n",
    "        [\n",
    "            A.Resize(height=512, width=512, p=1.0),\n",
    "            ToTensorV2(p=1.0),\n",
    "        ], \n",
    "        p=1.0, \n",
    "        bbox_params=A.BboxParams(\n",
    "            format='pascal_voc',\n",
    "            min_area=0, \n",
    "            min_visibility=0,\n",
    "            label_fields=['labels']\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_ROOT_PATH = '/home/hy/dataset/gwd/train'\n",
    "\n",
    "class DatasetRetriever(Dataset):\n",
    "\n",
    "    def __init__(self, marking, image_ids, transforms=None, test=False):\n",
    "        super().__init__()\n",
    "\n",
    "        self.image_ids = image_ids\n",
    "        self.marking = marking\n",
    "        self.transforms = transforms\n",
    "        self.test = test\n",
    "\n",
    "    def __getitem__(self, index: int):\n",
    "        image_id = self.image_ids[index]\n",
    "        \n",
    "        if self.test or random.random() > 0.5:\n",
    "            image, boxes = self.load_image_and_boxes(index)\n",
    "        else:\n",
    "            image, boxes = self.load_cutmix_image_and_boxes(index)\n",
    "\n",
    "        # there is only one class\n",
    "        labels = torch.ones((boxes.shape[0],), dtype=torch.int64)\n",
    "        \n",
    "        target = {}\n",
    "        target['boxes'] = boxes\n",
    "        target['labels'] = labels\n",
    "        target['image_id'] = torch.tensor([index])\n",
    "\n",
    "        if self.transforms:\n",
    "            for i in range(10):\n",
    "                sample = self.transforms(**{\n",
    "                    'image': image,\n",
    "                    'bboxes': target['boxes'],\n",
    "                    'labels': labels\n",
    "                })\n",
    "                if len(sample['bboxes']) > 0:\n",
    "                    image = sample['image']\n",
    "                    target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n",
    "                    target['boxes'][:,[0,1,2,3]] = target['boxes'][:,[1,0,3,2]]  #yxyx: be warning\n",
    "                    break\n",
    "\n",
    "        return image, target, image_id\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return self.image_ids.shape[0]\n",
    "\n",
    "    def load_image_and_boxes(self, index):\n",
    "        image_id = self.image_ids[index]\n",
    "        image = cv2.imread(f'{TRAIN_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n",
    "        image /= 255.0\n",
    "        records = self.marking[self.marking['image_id'] == image_id]\n",
    "        boxes = records[['x', 'y', 'w', 'h']].values\n",
    "        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n",
    "        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n",
    "        return image, boxes\n",
    "\n",
    "    def load_cutmix_image_and_boxes(self, index, imsize=1024):\n",
    "        \"\"\" \n",
    "        This implementation of cutmix author:  https://www.kaggle.com/nvnnghia \n",
    "        Refactoring and adaptation: https://www.kaggle.com/shonenkov\n",
    "        \"\"\"\n",
    "        w, h = imsize, imsize\n",
    "        s = imsize // 2\n",
    "    \n",
    "        xc, yc = [int(random.uniform(imsize * 0.25, imsize * 0.75)) for _ in range(2)]  # center x, y\n",
    "        indexes = [index] + [random.randint(0, self.image_ids.shape[0] - 1) for _ in range(3)]\n",
    "\n",
    "        result_image = np.full((imsize, imsize, 3), 1, dtype=np.float32)\n",
    "        result_boxes = []\n",
    "\n",
    "        for i, index in enumerate(indexes):\n",
    "            image, boxes = self.load_image_and_boxes(index)\n",
    "            if i == 0:\n",
    "                x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)\n",
    "                x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)\n",
    "            elif i == 1:  # top right\n",
    "                x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc\n",
    "                x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h\n",
    "            elif i == 2:  # bottom left\n",
    "                x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)\n",
    "                x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, max(xc, w), min(y2a - y1a, h)\n",
    "            elif i == 3:  # bottom right\n",
    "                x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)\n",
    "                x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)\n",
    "            result_image[y1a:y2a, x1a:x2a] = image[y1b:y2b, x1b:x2b]\n",
    "            padw = x1a - x1b\n",
    "            padh = y1a - y1b\n",
    "\n",
    "            boxes[:, 0] += padw\n",
    "            boxes[:, 1] += padh\n",
    "            boxes[:, 2] += padw\n",
    "            boxes[:, 3] += padh\n",
    "\n",
    "            result_boxes.append(boxes)\n",
    "\n",
    "        result_boxes = np.concatenate(result_boxes, 0)\n",
    "        np.clip(result_boxes[:, 0:], 0, 2 * s, out=result_boxes[:, 0:])\n",
    "        result_boxes = result_boxes.astype(np.int32)\n",
    "        result_boxes = result_boxes[np.where((result_boxes[:,2]-result_boxes[:,0])*(result_boxes[:,3]-result_boxes[:,1]) > 0)]\n",
    "        return result_image, result_boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fold_number = 0\n",
    "\n",
    "train_dataset = DatasetRetriever(\n",
    "    image_ids=df_folds[df_folds['fold'] != fold_number].index.values,\n",
    "    marking=marking,\n",
    "    transforms=get_train_transforms(),\n",
    "    test=False,\n",
    ")\n",
    "\n",
    "validation_dataset = DatasetRetriever(\n",
    "    image_ids=df_folds[df_folds['fold'] == fold_number].index.values,\n",
    "    marking=marking,\n",
    "    transforms=get_valid_transforms(),\n",
    "    test=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image, target, image_id = train_dataset[1]\n",
    "boxes = target['boxes'].cpu().numpy().astype(np.int32)\n",
    "\n",
    "numpy_image = image.permute(1,2,0).cpu().numpy()\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n",
    "\n",
    "for box in boxes:\n",
    "    cv2.rectangle(numpy_image, (box[1], box[0]), (box[3],  box[2]), (0, 1, 0), 2)\n",
    "    \n",
    "ax.set_axis_off()\n",
    "ax.imshow(numpy_image);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AverageMeter(object):\n",
    "    \"\"\"Computes and stores the average and current value\"\"\"\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "import torch.nn as nn\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "class Fitter:\n",
    "    \n",
    "    def __init__(self, model, device, config):\n",
    "        self.config = config\n",
    "        self.epoch = 0\n",
    "\n",
    "        self.base_dir = f'./{config.folder}'\n",
    "        if not os.path.exists(self.base_dir):\n",
    "            os.makedirs(self.base_dir)\n",
    "        \n",
    "        self.log_path = f'{self.base_dir}/log.txt'\n",
    "        self.best_summary_loss = 10**5\n",
    "        self.model = model\n",
    "        #self.model = nn.DataParallel(model)\n",
    "        self.device = device\n",
    "\n",
    "        param_optimizer = list(self.model.named_parameters())\n",
    "        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n",
    "        optimizer_grouped_parameters = [\n",
    "            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},\n",
    "            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n",
    "        ] \n",
    "\n",
    "        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=config.lr)\n",
    "        #self.optimizer = QHAdamW(self.model.parameters(), lr=config.lr)\n",
    "        self.scheduler = config.SchedulerClass(self.optimizer, **config.scheduler_params)\n",
    "        self.log(f'Fitter prepared. Device is {self.device}')\n",
    "\n",
    "    def fit(self, train_loader, validation_loader):\n",
    "        for e in range(self.config.n_epochs):\n",
    "            if self.config.verbose:\n",
    "                lr = self.optimizer.param_groups[0]['lr']\n",
    "                timestamp = datetime.utcnow().isoformat()\n",
    "                self.log(f'\\n{timestamp}\\nLR: {lr}')\n",
    "\n",
    "            t = time.time()\n",
    "            summary_loss = self.train_one_epoch(train_loader)\n",
    "\n",
    "            self.log(f'[RESULT]: Train. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')\n",
    "            self.save(f'{self.base_dir}/last-checkpoint.bin')\n",
    "\n",
    "            t = time.time()\n",
    "            summary_loss = self.validation(validation_loader)\n",
    "\n",
    "            self.log(f'[RESULT]: Val. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')\n",
    "            if summary_loss.avg < self.best_summary_loss:\n",
    "                self.best_summary_loss = summary_loss.avg\n",
    "                self.model.eval()\n",
    "                self.save(f'{self.base_dir}/best-checkpoint-{str(self.epoch).zfill(3)}epoch.bin')\n",
    "                for path in sorted(glob(f'{self.base_dir}/best-checkpoint-*epoch.bin'))[:-3]:\n",
    "                    os.remove(path)\n",
    "\n",
    "            if self.config.validation_scheduler:\n",
    "                self.scheduler.step(metrics=summary_loss.avg)\n",
    "\n",
    "            self.epoch += 1\n",
    "\n",
    "    def validation(self, val_loader):\n",
    "        self.model.eval()\n",
    "        summary_loss = AverageMeter()\n",
    "        t = time.time()\n",
    "        for step, (images, targets, image_ids) in enumerate(val_loader):\n",
    "            if self.config.verbose:\n",
    "                if step % self.config.verbose_step == 0:\n",
    "                    print(\n",
    "                        f'Val Step {step}/{len(val_loader)}, ' + \\\n",
    "                        f'summary_loss: {summary_loss.avg:.5f}, ' + \\\n",
    "                        f'time: {(time.time() - t):.5f}', end='\\r'\n",
    "                    )\n",
    "            with torch.no_grad():\n",
    "                images = torch.stack(images)\n",
    "                batch_size = images.shape[0]\n",
    "                images = images.to(self.device).float()\n",
    "                boxes = [target['boxes'].to(self.device).float() for target in targets]\n",
    "                labels = [target['labels'].to(self.device).float() for target in targets]\n",
    "\n",
    "                loss, _, _ = self.model(images, boxes, labels)\n",
    "                summary_loss.update(loss.detach().item(), batch_size)\n",
    "\n",
    "        return summary_loss\n",
    "\n",
    "    def train_one_epoch(self, train_loader):\n",
    "        self.model.train()\n",
    "        summary_loss = AverageMeter()\n",
    "        t = time.time()\n",
    "        for step, (images, targets, image_ids) in enumerate(train_loader):\n",
    "            if self.config.verbose:\n",
    "                if step % self.config.verbose_step == 0:\n",
    "                    print(\n",
    "                        f'Train Step {step}/{len(train_loader)}, ' + \\\n",
    "                        f'summary_loss: {summary_loss.avg:.5f}, ' + \\\n",
    "                        f'time: {(time.time() - t):.5f}', end='\\r'\n",
    "                    )\n",
    "            \n",
    "            images = torch.stack(images)\n",
    "            images = images.to(self.device).float()\n",
    "            batch_size = images.shape[0]\n",
    "            boxes = [target['boxes'].to(self.device).float() for target in targets]\n",
    "            labels = [target['labels'].to(self.device).float() for target in targets]\n",
    "\n",
    "            self.optimizer.zero_grad()\n",
    "            \n",
    "            loss, _, _ = self.model(images, boxes, labels)\n",
    "            \n",
    "            loss.backward()\n",
    "\n",
    "            summary_loss.update(loss.detach().item(), batch_size)\n",
    "\n",
    "            self.optimizer.step()\n",
    "\n",
    "            if self.config.step_scheduler:\n",
    "                self.scheduler.step()\n",
    "\n",
    "        return summary_loss\n",
    "    \n",
    "    def save(self, path):\n",
    "        self.model.eval()\n",
    "        torch.save({\n",
    "            'model_state_dict': self.model.model.state_dict(),\n",
    "            'optimizer_state_dict': self.optimizer.state_dict(),\n",
    "            'scheduler_state_dict': self.scheduler.state_dict(),\n",
    "            'best_summary_loss': self.best_summary_loss,\n",
    "            'epoch': self.epoch,\n",
    "        }, path)\n",
    "\n",
    "    def load(self, path):\n",
    "        checkpoint = torch.load(path)\n",
    "        self.model.model.load_state_dict(checkpoint['model_state_dict'])\n",
    "        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n",
    "        self.best_summary_loss = checkpoint['best_summary_loss']\n",
    "        self.epoch = checkpoint['epoch'] + 1\n",
    "        \n",
    "    def log(self, message):\n",
    "        if self.config.verbose:\n",
    "            print(message)\n",
    "        with open(self.log_path, 'a+') as logger:\n",
    "            logger.write(f'{message}\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainGlobalConfig:\n",
    "    num_workers = 4\n",
    "    batch_size = 8\n",
    "    #n_epochs = 40\n",
    "    n_epochs = 50\n",
    "    lr = 0.0002\n",
    "    \n",
    "    folder = '0525_effdet5-cutmix-augmix_test'\n",
    "\n",
    "    # -------------------\n",
    "    verbose = True\n",
    "    verbose_step = 1\n",
    "    # -------------------\n",
    "\n",
    "    # --------------------\n",
    "    step_scheduler = False  # do scheduler.step after optimizer.step\n",
    "    validation_scheduler = True  # do scheduler.step after validation stage loss\n",
    "\n",
    "#     SchedulerClass = torch.optim.lr_scheduler.OneCycleLR\n",
    "#     scheduler_params = dict(\n",
    "#         max_lr=0.001,\n",
    "#         epochs=n_epochs,\n",
    "#         steps_per_epoch=int(len(train_dataset) / batch_size),\n",
    "#         pct_start=0.1,\n",
    "#         anneal_strategy='cos', \n",
    "#         final_div_factor=10**5\n",
    "#     )\n",
    "    \n",
    "    SchedulerClass = torch.optim.lr_scheduler.ReduceLROnPlateau\n",
    "    scheduler_params = dict(\n",
    "        mode='min',\n",
    "        factor=0.5,\n",
    "        patience=1,\n",
    "        verbose=False, \n",
    "        threshold=0.0001,\n",
    "        threshold_mode='abs',\n",
    "        cooldown=0, \n",
    "        min_lr=1e-8,\n",
    "        eps=1e-08\n",
    "    )\n",
    "    # --------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(batch):\n",
    "    return tuple(zip(*batch))\n",
    "\n",
    "def run_training():\n",
    "    device = torch.device('cuda:0')\n",
    "    net.to(device)\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        train_dataset,\n",
    "        batch_size=TrainGlobalConfig.batch_size,\n",
    "        sampler=RandomSampler(train_dataset),\n",
    "        pin_memory=False,\n",
    "        drop_last=True,\n",
    "        num_workers=TrainGlobalConfig.num_workers,\n",
    "        collate_fn=collate_fn,\n",
    "    )\n",
    "    val_loader = torch.utils.data.DataLoader(\n",
    "        validation_dataset, \n",
    "        batch_size=TrainGlobalConfig.batch_size,\n",
    "        num_workers=TrainGlobalConfig.num_workers,\n",
    "        shuffle=False,\n",
    "        sampler=SequentialSampler(validation_dataset),\n",
    "        pin_memory=False,\n",
    "        collate_fn=collate_fn,\n",
    "    )\n",
    "\n",
    "    fitter = Fitter(model=net, device=device, config=TrainGlobalConfig)\n",
    "    fitter.fit(train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from effdet import get_efficientdet_config, EfficientDet, DetBenchTrain\n",
    "from effdet.efficientdet import HeadNet\n",
    "\n",
    "def get_net():\n",
    "    config = get_efficientdet_config('tf_efficientdet_d5')\n",
    "    net = EfficientDet(config, pretrained_backbone=False)\n",
    "    checkpoint = torch.load('efficientdet_d5-ef44aea8.pth')\n",
    "    net.load_state_dict(checkpoint)\n",
    "    config.num_classes = 1\n",
    "    config.image_size = 512\n",
    "    net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))\n",
    "    return DetBenchTrain(net, config)\n",
    "\n",
    "net = get_net()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_training()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Thank you for reading my kernel!\n",
    "\n",
    "So, I have prepared good training SOTA-model baseline for you, my friends! I have used n_epochs = 40 and have got best checkpoint single model that gives 0.7176 LB. You can see [here](https://www.kaggle.com/shonenkov/inference-efficientdet) inference kernel.\n",
    "\n",
    "Just recently I have started publishing my works, if you like this format of notebooks I would like continue to make kernels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:badeda]",
   "language": "python",
   "name": "conda-env-badeda-py"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}