pytorch-metric-learning-pipeline-only-images.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.016546,
     "end_time": "2021-03-30T17:09:02.125180",
     "exception": false,
     "start_time": "2021-03-30T17:09:02.108634",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Update\n",
    "\n",
    "Latest Version : Updated the LR scheduler with a custom Step LR schedule similar to what Chris Deotte used in Melanoma Competition. The intuition being the noise present in the labels\n",
    "\n",
    "Next Update : A better cross validation to simulate the test data\n",
    "\n",
    "# About this Notebook\n",
    "\n",
    "Ragnar [here](https://www.kaggle.com/ragnar123/shopee-efficientnetb3-arcmarginproduct) has showed us how to train efficientb3 with arcface and CrossEntropy Loss\n",
    "\n",
    "With this notebook here you can train any efficientnet b0-b7 with following three Metric Learning techniques above the Cross-Entropy Loss:\n",
    "* ArcFace : Most Popular in this competition \n",
    "* CosFace : https://arxiv.org/abs/1801.09414\n",
    "* Adacos : https://arxiv.org/abs/1905.00292\n",
    "\n",
    "Don't worry if these terms feel alien to you , following are few resources to understand them:\n",
    "* https://www.kaggle.com/c/shopee-product-matching/discussion/226279 --- Beautiful Explanantion by Chris Deotte\n",
    "* https://www.kaggle.com/slawekbiel/arcface-explained --  Beautiful kernel explaining ArcFace\n",
    "\n",
    "\n",
    "The training strategy used is as follows:\n",
    "* Make 5-folds stratified on label groups\n",
    "* Use one of the metric learning losses to predict the label  groups using cross entropy loss\n",
    "\n",
    "This notebook basically just converts ragnar's notebook into your beloved Pytorch . \n",
    "<font color='red'>Inference notebook for this training notebook can be found </font> [here](https://www.kaggle.com/tanulsingh077/metric-learning-image-tfidf-inference?scriptVersionId=57597192)\n",
    "\n",
    "<br>I have quick run this notebook because I have exhausted my gpu quota for this week . To see how this pipeline runs you can fork and run it .Hope this helps\n",
    "\n",
    "### Techniques planned to be added\n",
    "\n",
    "* Multiface posted by Mobassir [here](https://www.kaggle.com/c/shopee-product-matching/discussion/227383)\n",
    "\n",
    "You can also have a look at Siamese Type Training Example [here](https://www.kaggle.com/tanulsingh077/siamese-style-training-efficient-net-b0-on-tpu-s)\n",
    "\n",
    "\n",
    "# Important Note\n",
    "\n",
    "Please note that you can easily plug and play with this notebook, just by changing the values in the configuration. I have yet to analyze the data , but a lot of people have reported in noise in the dataset and thus using good augmentations like mixup,cutmix might help with training . Also losses like label smoothing , Focal loss,etc can be used. \n",
    "\n",
    "All my ideas I will keep on appending in this kernel and all the results will be maintained in [this](https://www.kaggle.com/c/shopee-product-matching/discussion/228148) discussion thread"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:02.160115Z",
     "iopub.status.busy": "2021-03-30T17:09:02.159497Z",
     "iopub.status.idle": "2021-03-30T17:09:02.162528Z",
     "shell.execute_reply": "2021-03-30T17:09:02.162003Z"
    },
    "papermill": {
     "duration": 0.022241,
     "end_time": "2021-03-30T17:09:02.162679",
     "exception": false,
     "start_time": "2021-03-30T17:09:02.140438",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../input/pytorch-image-models/pytorch-image-models-master')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:02.199448Z",
     "iopub.status.busy": "2021-03-30T17:09:02.198945Z",
     "iopub.status.idle": "2021-03-30T17:09:06.242263Z",
     "shell.execute_reply": "2021-03-30T17:09:06.241240Z"
    },
    "papermill": {
     "duration": 4.064166,
     "end_time": "2021-03-30T17:09:06.242395",
     "exception": false,
     "start_time": "2021-03-30T17:09:02.178229",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Preliminaries\n",
    "from tqdm import tqdm\n",
    "import math\n",
    "import random\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "# Visuals and CV2\n",
    "import cv2\n",
    "\n",
    "# albumentations for augs\n",
    "import albumentations\n",
    "from albumentations.pytorch.transforms import ToTensorV2\n",
    "\n",
    "#torch\n",
    "import torch\n",
    "import timm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Parameter\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import Dataset,DataLoader\n",
    "from torch.optim import Adam\n",
    "from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\n",
    "from torch.optim import Adam, lr_scheduler\n",
    "from torch.optim.lr_scheduler import _LRScheduler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.015447,
     "end_time": "2021-03-30T17:09:06.273946",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.258499",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.314469Z",
     "iopub.status.busy": "2021-03-30T17:09:06.313305Z",
     "iopub.status.idle": "2021-03-30T17:09:06.316135Z",
     "shell.execute_reply": "2021-03-30T17:09:06.315712Z"
    },
    "papermill": {
     "duration": 0.026836,
     "end_time": "2021-03-30T17:09:06.316240",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.289404",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "DIM = (512,512)\n",
    "\n",
    "NUM_WORKERS = 4\n",
    "TRAIN_BATCH_SIZE = 16\n",
    "VALID_BATCH_SIZE = 16\n",
    "EPOCHS = 5\n",
    "SEED = 42\n",
    "#LR = 3e-4\n",
    "\n",
    "device = torch.device('cuda')\n",
    "\n",
    "\n",
    "################################################# MODEL ####################################################################\n",
    "\n",
    "model_name = 'efficientnet_b3' #efficientnet_b0-b7\n",
    "\n",
    "################################################ Metric Loss and its params #######################################################\n",
    "loss_module = 'arcface' #'cosface' #'adacos'\n",
    "s = 30.0\n",
    "m = 0.5 \n",
    "ls_eps = 0.0\n",
    "easy_margin = False\n",
    "\n",
    "\n",
    "####################################### Scheduler and its params ############################################################\n",
    "# SCHEDULER = 'CosineAnnealingWarmRestarts' #'CosineAnnealingLR'\n",
    "# factor=0.2 # ReduceLROnPlateau\n",
    "# patience=4 # ReduceLROnPlateau\n",
    "# eps=1e-6 # ReduceLROnPlateau\n",
    "# T_max=10 # CosineAnnealingLR\n",
    "# T_0=4 # CosineAnnealingWarmRestarts\n",
    "# min_lr=1e-6\n",
    "\n",
    "\n",
    "scheduler_params = {\n",
    "        \"lr_start\": 1e-5,\n",
    "        \"lr_max\": 1e-5 * TRAIN_BATCH_SIZE,\n",
    "        \"lr_min\": 1e-6,\n",
    "        \"lr_ramp_ep\": 5,\n",
    "        \"lr_sus_ep\": 0,\n",
    "        \"lr_decay\": 0.8,\n",
    "    }\n",
    "\n",
    "############################################## Model Params ###############################################################\n",
    "model_params = {\n",
    "    'n_classes':11014,\n",
    "    'model_name':'efficientnet_b3',\n",
    "    'use_fc':False,\n",
    "    'fc_dim':512,\n",
    "    'dropout':0.0,\n",
    "    'loss_module':loss_module,\n",
    "    's':30.0,\n",
    "    'margin':0.50,\n",
    "    'ls_eps':0.0,\n",
    "    'theta_zero':0.785,\n",
    "    'pretrained':True\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.015218,
     "end_time": "2021-03-30T17:09:06.346817",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.331599",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.383377Z",
     "iopub.status.busy": "2021-03-30T17:09:06.382838Z",
     "iopub.status.idle": "2021-03-30T17:09:06.388992Z",
     "shell.execute_reply": "2021-03-30T17:09:06.388460Z"
    },
    "papermill": {
     "duration": 0.025718,
     "end_time": "2021-03-30T17:09:06.389095",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.363377",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def seed_torch(seed=42):\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",
    "    \n",
    "seed_torch(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.425194Z",
     "iopub.status.busy": "2021-03-30T17:09:06.424643Z",
     "iopub.status.idle": "2021-03-30T17:09:06.428059Z",
     "shell.execute_reply": "2021-03-30T17:09:06.428419Z"
    },
    "papermill": {
     "duration": 0.023798,
     "end_time": "2021-03-30T17:09:06.428535",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.404737",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class AverageMeter(object):\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": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.465652Z",
     "iopub.status.busy": "2021-03-30T17:09:06.465020Z",
     "iopub.status.idle": "2021-03-30T17:09:06.467853Z",
     "shell.execute_reply": "2021-03-30T17:09:06.467446Z"
    },
    "papermill": {
     "duration": 0.023302,
     "end_time": "2021-03-30T17:09:06.467966",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.444664",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def fetch_scheduler(optimizer):\n",
    "        if SCHEDULER =='ReduceLROnPlateau':\n",
    "            scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=factor, patience=patience, verbose=True, eps=eps)\n",
    "        elif SCHEDULER =='CosineAnnealingLR':\n",
    "            scheduler = CosineAnnealingLR(optimizer, T_max=T_max, eta_min=min_lr, last_epoch=-1)\n",
    "        elif SCHEDULER =='CosineAnnealingWarmRestarts':\n",
    "            scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=T_0, T_mult=1, eta_min=min_lr, last_epoch=-1)\n",
    "        return scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.502894Z",
     "iopub.status.busy": "2021-03-30T17:09:06.502285Z",
     "iopub.status.idle": "2021-03-30T17:09:06.505184Z",
     "shell.execute_reply": "2021-03-30T17:09:06.504742Z"
    },
    "papermill": {
     "duration": 0.021543,
     "end_time": "2021-03-30T17:09:06.505288",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.483745",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def fetch_loss():\n",
    "    loss = nn.CrossEntropyLoss()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.015666,
     "end_time": "2021-03-30T17:09:06.536713",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.521047",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Augmentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.574995Z",
     "iopub.status.busy": "2021-03-30T17:09:06.574306Z",
     "iopub.status.idle": "2021-03-30T17:09:06.577008Z",
     "shell.execute_reply": "2021-03-30T17:09:06.576604Z"
    },
    "papermill": {
     "duration": 0.024564,
     "end_time": "2021-03-30T17:09:06.577110",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.552546",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_train_transforms():\n",
    "    return albumentations.Compose(\n",
    "        [   \n",
    "            albumentations.Resize(DIM[0],DIM[1],always_apply=True),\n",
    "            albumentations.HorizontalFlip(p=0.5),\n",
    "            albumentations.VerticalFlip(p=0.5),\n",
    "            albumentations.Rotate(limit=120, p=0.8),\n",
    "            albumentations.RandomBrightness(limit=(0.09, 0.6), p=0.5),\n",
    "            #albumentations.Cutout(num_holes=8, max_h_size=8, max_w_size=8, fill_value=0, always_apply=False, p=0.5),\n",
    "            #albumentations.ShiftScaleRotate(\n",
    "              #  shift_limit=0.25, scale_limit=0.1, rotate_limit=0\n",
    "            #),\n",
    "            albumentations.Normalize(),\n",
    "            ToTensorV2(p=1.0),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "def get_valid_transforms():\n",
    "\n",
    "    return albumentations.Compose(\n",
    "        [\n",
    "            albumentations.Resize(DIM[0],DIM[1],always_apply=True),\n",
    "            albumentations.Normalize(),\n",
    "        ToTensorV2(p=1.0)\n",
    "        ]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.015715,
     "end_time": "2021-03-30T17:09:06.608608",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.592893",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.646523Z",
     "iopub.status.busy": "2021-03-30T17:09:06.645881Z",
     "iopub.status.idle": "2021-03-30T17:09:06.648698Z",
     "shell.execute_reply": "2021-03-30T17:09:06.648319Z"
    },
    "papermill": {
     "duration": 0.024233,
     "end_time": "2021-03-30T17:09:06.648796",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.624563",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ShopeeDataset(Dataset):\n",
    "    def __init__(self, csv, transforms=None):\n",
    "\n",
    "        self.csv = csv.reset_index()\n",
    "        self.augmentations = transforms\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.csv.shape[0]\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        row = self.csv.iloc[index]\n",
    "        \n",
    "        text = row.title\n",
    "        \n",
    "        image = cv2.imread(row.filepath)\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "        \n",
    "        if self.augmentations:\n",
    "            augmented = self.augmentations(image=image)\n",
    "            image = augmented['image']       \n",
    "        \n",
    "        \n",
    "        return image,torch.tensor(row.label_group)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.015947,
     "end_time": "2021-03-30T17:09:06.680976",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.665029",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.727718Z",
     "iopub.status.busy": "2021-03-30T17:09:06.727027Z",
     "iopub.status.idle": "2021-03-30T17:09:06.729836Z",
     "shell.execute_reply": "2021-03-30T17:09:06.729355Z"
    },
    "papermill": {
     "duration": 0.032737,
     "end_time": "2021-03-30T17:09:06.729950",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.697213",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ShopeeNet(nn.Module):\n",
    "\n",
    "    def __init__(self,\n",
    "                 n_classes,\n",
    "                 model_name='efficientnet_b0',\n",
    "                 use_fc=False,\n",
    "                 fc_dim=512,\n",
    "                 dropout=0.0,\n",
    "                 loss_module='softmax',\n",
    "                 s=30.0,\n",
    "                 margin=0.50,\n",
    "                 ls_eps=0.0,\n",
    "                 theta_zero=0.785,\n",
    "                 pretrained=True):\n",
    "        \"\"\"\n",
    "        :param n_classes:\n",
    "        :param model_name: name of model from pretrainedmodels\n",
    "            e.g. resnet50, resnext101_32x4d, pnasnet5large\n",
    "        :param pooling: One of ('SPoC', 'MAC', 'RMAC', 'GeM', 'Rpool', 'Flatten', 'CompactBilinearPooling')\n",
    "        :param loss_module: One of ('arcface', 'cosface', 'softmax')\n",
    "        \"\"\"\n",
    "        super(ShopeeNet, self).__init__()\n",
    "        print('Building Model Backbone for {} model'.format(model_name))\n",
    "\n",
    "        self.backbone = timm.create_model(model_name, pretrained=pretrained)\n",
    "        final_in_features = self.backbone.classifier.in_features\n",
    "        \n",
    "        self.backbone.classifier = nn.Identity()\n",
    "        self.backbone.global_pool = nn.Identity()\n",
    "        \n",
    "        self.pooling =  nn.AdaptiveAvgPool2d(1)\n",
    "            \n",
    "        self.use_fc = use_fc\n",
    "        if use_fc:\n",
    "            self.dropout = nn.Dropout(p=dropout)\n",
    "            self.fc = nn.Linear(final_in_features, fc_dim)\n",
    "            self.bn = nn.BatchNorm1d(fc_dim)\n",
    "            self._init_params()\n",
    "            final_in_features = fc_dim\n",
    "\n",
    "        self.loss_module = loss_module\n",
    "        if loss_module == 'arcface':\n",
    "            self.final = ArcMarginProduct(final_in_features, n_classes,\n",
    "                                          s=s, m=margin, easy_margin=False, ls_eps=ls_eps)\n",
    "        elif loss_module == 'cosface':\n",
    "            self.final = AddMarginProduct(final_in_features, n_classes, s=s, m=margin)\n",
    "        elif loss_module == 'adacos':\n",
    "            self.final = AdaCos(final_in_features, n_classes, m=margin, theta_zero=theta_zero)\n",
    "        else:\n",
    "            self.final = nn.Linear(final_in_features, n_classes)\n",
    "\n",
    "    def _init_params(self):\n",
    "        nn.init.xavier_normal_(self.fc.weight)\n",
    "        nn.init.constant_(self.fc.bias, 0)\n",
    "        nn.init.constant_(self.bn.weight, 1)\n",
    "        nn.init.constant_(self.bn.bias, 0)\n",
    "\n",
    "    def forward(self, x, label):\n",
    "        feature = self.extract_feat(x)\n",
    "        if self.loss_module in ('arcface', 'cosface', 'adacos'):\n",
    "            logits = self.final(feature, label)\n",
    "        else:\n",
    "            logits = self.final(feature)\n",
    "        return logits\n",
    "\n",
    "    def extract_feat(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        x = self.backbone(x)\n",
    "        x = self.pooling(x).view(batch_size, -1)\n",
    "\n",
    "        if self.use_fc:\n",
    "            x = self.dropout(x)\n",
    "            x = self.fc(x)\n",
    "            x = self.bn(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.016651,
     "end_time": "2021-03-30T17:09:06.762716",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.746065",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Metric Learning Losses\n",
    "\n",
    "* https://github.com/lyakaap/Landmark2019-1st-and-3rd-Place-Solution/blob/master/src/modeling/metric_learning.py -- Code Taken from here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.806125Z",
     "iopub.status.busy": "2021-03-30T17:09:06.805446Z",
     "iopub.status.idle": "2021-03-30T17:09:06.808200Z",
     "shell.execute_reply": "2021-03-30T17:09:06.807794Z"
    },
    "papermill": {
     "duration": 0.029062,
     "end_time": "2021-03-30T17:09:06.808327",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.779265",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class AdaCos(nn.Module):\n",
    "    def __init__(self, in_features, out_features, m=0.50, ls_eps=0, theta_zero=math.pi/4):\n",
    "        super(AdaCos, self).__init__()\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.theta_zero = theta_zero\n",
    "        self.s = math.log(out_features - 1) / math.cos(theta_zero)\n",
    "        self.m = m\n",
    "        self.ls_eps = ls_eps  # label smoothing\n",
    "        self.weight = Parameter(torch.FloatTensor(out_features, in_features))\n",
    "        nn.init.xavier_uniform_(self.weight)\n",
    "\n",
    "    def forward(self, input, label):\n",
    "        # normalize features\n",
    "        x = F.normalize(input)\n",
    "        # normalize weights\n",
    "        W = F.normalize(self.weight)\n",
    "        # dot product\n",
    "        logits = F.linear(x, W)\n",
    "        # add margin\n",
    "        theta = torch.acos(torch.clamp(logits, -1.0 + 1e-7, 1.0 - 1e-7))\n",
    "        target_logits = torch.cos(theta + self.m)\n",
    "        one_hot = torch.zeros_like(logits)\n",
    "        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n",
    "        if self.ls_eps > 0:\n",
    "            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features\n",
    "        output = logits * (1 - one_hot) + target_logits * one_hot\n",
    "        # feature re-scale\n",
    "        with torch.no_grad():\n",
    "            B_avg = torch.where(one_hot < 1, torch.exp(self.s * logits), torch.zeros_like(logits))\n",
    "            B_avg = torch.sum(B_avg) / input.size(0)\n",
    "            theta_med = torch.median(theta)\n",
    "            self.s = torch.log(B_avg) / torch.cos(torch.min(self.theta_zero * torch.ones_like(theta_med), theta_med))\n",
    "        output *= self.s\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.853934Z",
     "iopub.status.busy": "2021-03-30T17:09:06.853266Z",
     "iopub.status.idle": "2021-03-30T17:09:06.856292Z",
     "shell.execute_reply": "2021-03-30T17:09:06.855756Z"
    },
    "papermill": {
     "duration": 0.031689,
     "end_time": "2021-03-30T17:09:06.856391",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.824702",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ArcMarginProduct(nn.Module):\n",
    "    r\"\"\"Implement of large margin arc distance: :\n",
    "        Args:\n",
    "            in_features: size of each input sample\n",
    "            out_features: size of each output sample\n",
    "            s: norm of input feature\n",
    "            m: margin\n",
    "            cos(theta + m)\n",
    "        \"\"\"\n",
    "    def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False, ls_eps=0.0):\n",
    "        super(ArcMarginProduct, self).__init__()\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.s = s\n",
    "        self.m = m\n",
    "        self.ls_eps = ls_eps  # label smoothing\n",
    "        self.weight = Parameter(torch.FloatTensor(out_features, in_features))\n",
    "        nn.init.xavier_uniform_(self.weight)\n",
    "\n",
    "        self.easy_margin = easy_margin\n",
    "        self.cos_m = math.cos(m)\n",
    "        self.sin_m = math.sin(m)\n",
    "        self.th = math.cos(math.pi - m)\n",
    "        self.mm = math.sin(math.pi - m) * m\n",
    "\n",
    "    def forward(self, input, label):\n",
    "        # --------------------------- cos(theta) & phi(theta) ---------------------------\n",
    "        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n",
    "        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))\n",
    "        phi = cosine * self.cos_m - sine * self.sin_m\n",
    "        if self.easy_margin:\n",
    "            phi = torch.where(cosine > 0, phi, cosine)\n",
    "        else:\n",
    "            phi = torch.where(cosine > self.th, phi, cosine - self.mm)\n",
    "        # --------------------------- convert label to one-hot ---------------------------\n",
    "        # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda')\n",
    "        one_hot = torch.zeros(cosine.size(), device='cuda')\n",
    "        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n",
    "        if self.ls_eps > 0:\n",
    "            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features\n",
    "        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------\n",
    "        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n",
    "        output *= self.s\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.896994Z",
     "iopub.status.busy": "2021-03-30T17:09:06.896356Z",
     "iopub.status.idle": "2021-03-30T17:09:06.899276Z",
     "shell.execute_reply": "2021-03-30T17:09:06.898886Z"
    },
    "papermill": {
     "duration": 0.026357,
     "end_time": "2021-03-30T17:09:06.899376",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.873019",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class AddMarginProduct(nn.Module):\n",
    "    r\"\"\"Implement of large margin cosine distance: :\n",
    "    Args:\n",
    "        in_features: size of each input sample\n",
    "        out_features: size of each output sample\n",
    "        s: norm of input feature\n",
    "        m: margin\n",
    "        cos(theta) - m\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_features, out_features, s=30.0, m=0.40):\n",
    "        super(AddMarginProduct, self).__init__()\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.s = s\n",
    "        self.m = m\n",
    "        self.weight = Parameter(torch.FloatTensor(out_features, in_features))\n",
    "        nn.init.xavier_uniform_(self.weight)\n",
    "\n",
    "    def forward(self, input, label):\n",
    "        # --------------------------- cos(theta) & phi(theta) ---------------------------\n",
    "        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n",
    "        phi = cosine - self.m\n",
    "        # --------------------------- convert label to one-hot ---------------------------\n",
    "        one_hot = torch.zeros(cosine.size(), device='cuda')\n",
    "        # one_hot = one_hot.cuda() if cosine.is_cuda else one_hot\n",
    "        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n",
    "        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------\n",
    "        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)  # you can use torch.where if your torch.__version__ is 0.4\n",
    "        output *= self.s\n",
    "        # print(output)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.016291,
     "end_time": "2021-03-30T17:09:06.932492",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.916201",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Custom LR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:06.975450Z",
     "iopub.status.busy": "2021-03-30T17:09:06.974785Z",
     "iopub.status.idle": "2021-03-30T17:09:06.977956Z",
     "shell.execute_reply": "2021-03-30T17:09:06.978541Z"
    },
    "papermill": {
     "duration": 0.029644,
     "end_time": "2021-03-30T17:09:06.978690",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.949046",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ShopeeScheduler(_LRScheduler):\n",
    "    def __init__(self, optimizer, lr_start=5e-6, lr_max=1e-5,\n",
    "                 lr_min=1e-6, lr_ramp_ep=5, lr_sus_ep=0, lr_decay=0.8,\n",
    "                 last_epoch=-1):\n",
    "        self.lr_start = lr_start\n",
    "        self.lr_max = lr_max\n",
    "        self.lr_min = lr_min\n",
    "        self.lr_ramp_ep = lr_ramp_ep\n",
    "        self.lr_sus_ep = lr_sus_ep\n",
    "        self.lr_decay = lr_decay\n",
    "        super(ShopeeScheduler, self).__init__(optimizer, last_epoch)\n",
    "        \n",
    "    def get_lr(self):\n",
    "        if not self._get_lr_called_within_step:\n",
    "            warnings.warn(\"To get the last learning rate computed by the scheduler, \"\n",
    "                          \"please use `get_last_lr()`.\", UserWarning)\n",
    "        \n",
    "        if self.last_epoch == 0:\n",
    "            self.last_epoch += 1\n",
    "            return [self.lr_start for _ in self.optimizer.param_groups]\n",
    "        \n",
    "        lr = self._compute_lr_from_epoch()\n",
    "        self.last_epoch += 1\n",
    "        \n",
    "        return [lr for _ in self.optimizer.param_groups]\n",
    "    \n",
    "    def _get_closed_form_lr(self):\n",
    "        return self.base_lrs\n",
    "    \n",
    "    def _compute_lr_from_epoch(self):\n",
    "        if self.last_epoch < self.lr_ramp_ep:\n",
    "            lr = ((self.lr_max - self.lr_start) / \n",
    "                  self.lr_ramp_ep * self.last_epoch + \n",
    "                  self.lr_start)\n",
    "        \n",
    "        elif self.last_epoch < self.lr_ramp_ep + self.lr_sus_ep:\n",
    "            lr = self.lr_max\n",
    "            \n",
    "        else:\n",
    "            lr = ((self.lr_max - self.lr_min) * self.lr_decay**\n",
    "                  (self.last_epoch - self.lr_ramp_ep - self.lr_sus_ep) + \n",
    "                  self.lr_min)\n",
    "        return lr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.017739,
     "end_time": "2021-03-30T17:09:07.014476",
     "exception": false,
     "start_time": "2021-03-30T17:09:06.996737",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Training Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:07.056214Z",
     "iopub.status.busy": "2021-03-30T17:09:07.055548Z",
     "iopub.status.idle": "2021-03-30T17:09:07.058471Z",
     "shell.execute_reply": "2021-03-30T17:09:07.057895Z"
    },
    "papermill": {
     "duration": 0.026584,
     "end_time": "2021-03-30T17:09:07.058583",
     "exception": false,
     "start_time": "2021-03-30T17:09:07.031999",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_fn(dataloader,model,criterion,optimizer,device,scheduler,epoch):\n",
    "    model.train()\n",
    "    loss_score = AverageMeter()\n",
    "    \n",
    "    tk0 = tqdm(enumerate(dataloader), total=len(dataloader))\n",
    "    for bi,d in tk0:\n",
    "        \n",
    "        batch_size = d[0].shape[0]\n",
    "\n",
    "\n",
    "        images = d[0]\n",
    "        targets = d[1]\n",
    "\n",
    "        images = images.to(device)\n",
    "        targets = targets.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        output = model(images,targets)\n",
    "        \n",
    "        loss = criterion(output,targets)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        loss_score.update(loss.detach().item(), batch_size)\n",
    "        tk0.set_postfix(Train_Loss=loss_score.avg,Epoch=epoch,LR=optimizer.param_groups[0]['lr'])\n",
    "        \n",
    "    if scheduler is not None:\n",
    "            scheduler.step()\n",
    "        \n",
    "    return loss_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.018147,
     "end_time": "2021-03-30T17:09:07.093715",
     "exception": false,
     "start_time": "2021-03-30T17:09:07.075568",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Evaluation Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:07.141046Z",
     "iopub.status.busy": "2021-03-30T17:09:07.140128Z",
     "iopub.status.idle": "2021-03-30T17:09:07.143254Z",
     "shell.execute_reply": "2021-03-30T17:09:07.142726Z"
    },
    "papermill": {
     "duration": 0.028188,
     "end_time": "2021-03-30T17:09:07.143383",
     "exception": false,
     "start_time": "2021-03-30T17:09:07.115195",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def eval_fn(data_loader,model,criterion,device):\n",
    "    \n",
    "    loss_score = AverageMeter()\n",
    "    \n",
    "    model.eval()\n",
    "    tk0 = tqdm(enumerate(data_loader), total=len(data_loader))\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        \n",
    "        for bi,d in tk0:\n",
    "            batch_size = d[0].size()[0]\n",
    "\n",
    "            image = d[0]\n",
    "            targets = d[1]\n",
    "\n",
    "            image = image.to(device)\n",
    "            targets = targets.to(device)\n",
    "\n",
    "            output = model(image,targets)\n",
    "\n",
    "            loss = criterion(output,targets)\n",
    "            \n",
    "            loss_score.update(loss.detach().item(), batch_size)\n",
    "            tk0.set_postfix(Eval_Loss=loss_score.avg)\n",
    "            \n",
    "    return loss_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.017344,
     "end_time": "2021-03-30T17:09:07.177668",
     "exception": false,
     "start_time": "2021-03-30T17:09:07.160324",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:07.221716Z",
     "iopub.status.busy": "2021-03-30T17:09:07.221251Z",
     "iopub.status.idle": "2021-03-30T17:09:07.456863Z",
     "shell.execute_reply": "2021-03-30T17:09:07.456367Z"
    },
    "papermill": {
     "duration": 0.262278,
     "end_time": "2021-03-30T17:09:07.457007",
     "exception": false,
     "start_time": "2021-03-30T17:09:07.194729",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = pd.read_csv('../input/shopee-folds/folds.csv')\n",
    "data['filepath'] = data['image'].apply(lambda x: os.path.join('../input/shopee-product-matching/', 'train_images', x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:07.503181Z",
     "iopub.status.busy": "2021-03-30T17:09:07.502602Z",
     "iopub.status.idle": "2021-03-30T17:09:07.512717Z",
     "shell.execute_reply": "2021-03-30T17:09:07.513166Z"
    },
    "papermill": {
     "duration": 0.038965,
     "end_time": "2021-03-30T17:09:07.513290",
     "exception": false,
     "start_time": "2021-03-30T17:09:07.474325",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>posting_id</th>\n",
       "      <th>image</th>\n",
       "      <th>image_phash</th>\n",
       "      <th>title</th>\n",
       "      <th>label_group</th>\n",
       "      <th>fold</th>\n",
       "      <th>filepath</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>train_129225211</td>\n",
       "      <td>0000a68812bc7e98c42888dfb1c07da0.jpg</td>\n",
       "      <td>94974f937d4c2433</td>\n",
       "      <td>Paper Bag Victoria Secret</td>\n",
       "      <td>249114794</td>\n",
       "      <td>0</td>\n",
       "      <td>../input/shopee-product-matching/train_images/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>train_3386243561</td>\n",
       "      <td>00039780dfc94d01db8676fe789ecd05.jpg</td>\n",
       "      <td>af3f9460c2838f0f</td>\n",
       "      <td>Double Tape 3M VHB 12 mm x 4,5 m ORIGINAL / DO...</td>\n",
       "      <td>2937985045</td>\n",
       "      <td>2</td>\n",
       "      <td>../input/shopee-product-matching/train_images/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>train_2288590299</td>\n",
       "      <td>000a190fdd715a2a36faed16e2c65df7.jpg</td>\n",
       "      <td>b94cb00ed3e50f78</td>\n",
       "      <td>Maling TTS Canned Pork Luncheon Meat 397 gr</td>\n",
       "      <td>2395904891</td>\n",
       "      <td>0</td>\n",
       "      <td>../input/shopee-product-matching/train_images/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>train_2406599165</td>\n",
       "      <td>00117e4fc239b1b641ff08340b429633.jpg</td>\n",
       "      <td>8514fc58eafea283</td>\n",
       "      <td>Daster Batik Lengan pendek - Motif Acak / Camp...</td>\n",
       "      <td>4093212188</td>\n",
       "      <td>1</td>\n",
       "      <td>../input/shopee-product-matching/train_images/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>train_3369186413</td>\n",
       "      <td>00136d1cf4edede0203f32f05f660588.jpg</td>\n",
       "      <td>a6f319f924ad708c</td>\n",
       "      <td>Nescafe \\xc3\\x89clair Latte 220ml</td>\n",
       "      <td>3648931069</td>\n",
       "      <td>3</td>\n",
       "      <td>../input/shopee-product-matching/train_images/...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         posting_id                                 image       image_phash  \\\n",
       "0   train_129225211  0000a68812bc7e98c42888dfb1c07da0.jpg  94974f937d4c2433   \n",
       "1  train_3386243561  00039780dfc94d01db8676fe789ecd05.jpg  af3f9460c2838f0f   \n",
       "2  train_2288590299  000a190fdd715a2a36faed16e2c65df7.jpg  b94cb00ed3e50f78   \n",
       "3  train_2406599165  00117e4fc239b1b641ff08340b429633.jpg  8514fc58eafea283   \n",
       "4  train_3369186413  00136d1cf4edede0203f32f05f660588.jpg  a6f319f924ad708c   \n",
       "\n",
       "                                               title  label_group  fold  \\\n",
       "0                          Paper Bag Victoria Secret    249114794     0   \n",
       "1  Double Tape 3M VHB 12 mm x 4,5 m ORIGINAL / DO...   2937985045     2   \n",
       "2        Maling TTS Canned Pork Luncheon Meat 397 gr   2395904891     0   \n",
       "3  Daster Batik Lengan pendek - Motif Acak / Camp...   4093212188     1   \n",
       "4                  Nescafe \\xc3\\x89clair Latte 220ml   3648931069     3   \n",
       "\n",
       "                                            filepath  \n",
       "0  ../input/shopee-product-matching/train_images/...  \n",
       "1  ../input/shopee-product-matching/train_images/...  \n",
       "2  ../input/shopee-product-matching/train_images/...  \n",
       "3  ../input/shopee-product-matching/train_images/...  \n",
       "4  ../input/shopee-product-matching/train_images/...  "
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:07.553863Z",
     "iopub.status.busy": "2021-03-30T17:09:07.553114Z",
     "iopub.status.idle": "2021-03-30T17:09:07.561800Z",
     "shell.execute_reply": "2021-03-30T17:09:07.561403Z"
    },
    "papermill": {
     "duration": 0.030364,
     "end_time": "2021-03-30T17:09:07.561922",
     "exception": false,
     "start_time": "2021-03-30T17:09:07.531558",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "encoder = LabelEncoder()\n",
    "data['label_group'] = encoder.fit_transform(data['label_group'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:07.607053Z",
     "iopub.status.busy": "2021-03-30T17:09:07.606464Z",
     "iopub.status.idle": "2021-03-30T17:09:07.610231Z",
     "shell.execute_reply": "2021-03-30T17:09:07.609827Z"
    },
    "papermill": {
     "duration": 0.030818,
     "end_time": "2021-03-30T17:09:07.610345",
     "exception": false,
     "start_time": "2021-03-30T17:09:07.579527",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def run():\n",
    "        \n",
    "    train = data[data['fold']!=0].reset_index(drop=True)\n",
    "    valid = data[data['fold']==0].reset_index(drop=True)\n",
    "    # Defining DataSet\n",
    "    train_dataset = ShopeeDataset(\n",
    "        csv=train,\n",
    "        transforms=get_train_transforms(),\n",
    "    )\n",
    "        \n",
    "    valid_dataset = ShopeeDataset(\n",
    "        csv=valid,\n",
    "        transforms=get_valid_transforms(),\n",
    "    )\n",
    "        \n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        train_dataset,\n",
    "        batch_size=TRAIN_BATCH_SIZE,\n",
    "        pin_memory=True,\n",
    "        drop_last=True,\n",
    "        num_workers=NUM_WORKERS\n",
    "    )\n",
    "    \n",
    "    valid_loader = torch.utils.data.DataLoader(\n",
    "        valid_dataset,\n",
    "        batch_size=VALID_BATCH_SIZE,\n",
    "        num_workers=NUM_WORKERS,\n",
    "        shuffle=False,\n",
    "        pin_memory=True,\n",
    "        drop_last=False,\n",
    "    )\n",
    "    \n",
    "    # Defining Device\n",
    "    device = torch.device(\"cuda\")\n",
    "    \n",
    "    # Defining Model for specific fold\n",
    "    model = ShopeeNet(**model_params)\n",
    "    model.to(device)\n",
    "    \n",
    "    #DEfining criterion\n",
    "    criterion = fetch_loss()\n",
    "    criterion.to(device)\n",
    "        \n",
    "    # Defining Optimizer with weight decay to params other than bias and layer norms\n",
    "#     param_optimizer = list(model.named_parameters())\n",
    "#     no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n",
    "#     optimizer_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",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=scheduler_params['lr_start'])\n",
    "    \n",
    "    #Defining LR SCheduler\n",
    "    scheduler = ShopeeScheduler(optimizer,**scheduler_params)\n",
    "        \n",
    "    # THE ENGINE LOOP\n",
    "    best_loss = 10000\n",
    "    for epoch in range(EPOCHS):\n",
    "        train_loss = train_fn(train_loader, model,criterion, optimizer, device,scheduler=scheduler,epoch=epoch)\n",
    "        \n",
    "        valid_loss = eval_fn(valid_loader, model, criterion,device)\n",
    "        \n",
    "        if valid_loss.avg < best_loss:\n",
    "            best_loss = valid_loss.avg\n",
    "            torch.save(model.state_dict(),f'model_{model_name}_IMG_SIZE_{DIM[0]}_{loss_module}.bin')\n",
    "            print('best model found for epoch {}'.format(epoch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-30T17:09:07.651138Z",
     "iopub.status.busy": "2021-03-30T17:09:07.650477Z",
     "iopub.status.idle": "2021-03-30T18:40:57.680922Z",
     "shell.execute_reply": "2021-03-30T18:40:57.640948Z"
    },
    "papermill": {
     "duration": 5510.052704,
     "end_time": "2021-03-30T18:40:57.681051",
     "exception": false,
     "start_time": "2021-03-30T17:09:07.628347",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building Model Backbone for efficientnet_b3 model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/efficientnet_b3_ra2-cf984f9c.pth\" to /root/.cache/torch/hub/checkpoints/efficientnet_b3_ra2-cf984f9c.pth\n",
      "100%|██████████| 1712/1712 [16:29<00:00,  1.73it/s, Epoch=0, LR=1e-5, Train_Loss=23.8]\n",
      "100%|██████████| 429/429 [01:59<00:00,  3.60it/s, Eval_Loss=23.6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best model found for epoch 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1712/1712 [16:21<00:00,  1.74it/s, Epoch=1, LR=7e-5, Train_Loss=21.1]\n",
      "100%|██████████| 429/429 [01:58<00:00,  3.62it/s, Eval_Loss=19.6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best model found for epoch 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1712/1712 [16:18<00:00,  1.75it/s, Epoch=2, LR=0.00013, Train_Loss=17.8]\n",
      "100%|██████████| 429/429 [01:58<00:00,  3.62it/s, Eval_Loss=16.6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best model found for epoch 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1712/1712 [16:17<00:00,  1.75it/s, Epoch=3, LR=0.000128, Train_Loss=13.5]\n",
      "100%|██████████| 429/429 [01:54<00:00,  3.74it/s, Eval_Loss=14.2]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best model found for epoch 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1712/1712 [16:20<00:00,  1.75it/s, Epoch=4, LR=8.24e-5, Train_Loss=9.57]\n",
      "100%|██████████| 429/429 [01:56<00:00,  3.70it/s, Eval_Loss=12.8]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best model found for epoch 4\n"
     ]
    }
   ],
   "source": [
    "run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "papermill": {
     "duration": 5.471599,
     "end_time": "2021-03-30T18:41:08.848018",
     "exception": false,
     "start_time": "2021-03-30T18:41:03.376419",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.9"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 5540.142232,
   "end_time": "2021-03-30T18:41:17.522502",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2021-03-30T17:08:57.380270",
   "version": "2.2.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}