src/training/hotel-id-cosface-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5tPFJ89V3BFT"
   },
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aEV7y8x-WEKj"
   },
   "source": [
    "This notebook is inteded to run on colab on preprocessed images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vW7S-8qm3FfU"
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_ABy2M3C3HEb"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/gdrive')\n",
    "%cd /gdrive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1JM4AM6M2_aj"
   },
   "outputs": [],
   "source": [
    "# !pip install efficientnet_pytorch\n",
    "!pip install git+https://github.com/rwightman/pytorch-image-models\n",
    "!pip install pytorch-metric-learning\n",
    "!pip install faiss-gpu\n",
    "!pip install imgaug -U\n",
    "!pip install albumentations -U"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MyC4gTwZ3MKJ"
   },
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "id": "u0Bz2ktn2_ap"
   },
   "outputs": [],
   "source": [
    "import numpy as np # linear algebra\n",
    "import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n",
    "import random\n",
    "import os\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tOszKuxt3PXn"
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score\n",
    "from sklearn.utils import class_weight\n",
    "from PIL import Image as pil_image\n",
    "from tqdm import tqdm\n",
    "import scipy\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import plotly.express as px\n",
    "import plotly.graph_objects as go"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uQE7wYFR3QxV"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import timm\n",
    "from timm.optim import Lookahead, RAdam\n",
    "from pytorch_metric_learning import miners, losses, samplers , distances, regularizers "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tirOg6jm3aIB"
   },
   "source": [
    "# Global"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DV7qHDuYGoJH"
   },
   "outputs": [],
   "source": [
    "IMG_SIZE = 512\n",
    "SEED = 42\n",
    "PROJECT_FOLDER = \"/gdrive/MyDrive/Projects/Hotel-ID/\"\n",
    "DATA_FOLDER = \"/home/data/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TB9CXg8U3bbQ"
   },
   "outputs": [],
   "source": [
    "!mkdir {DATA_FOLDER}\n",
    "!unzip -qq {PROJECT_FOLDER}data/train-{IMG_SIZE}x{IMG_SIZE}.zip -d /home/data/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1wH0zWUS2_aq"
   },
   "outputs": [],
   "source": [
    "print(os.listdir(PROJECT_FOLDER))\n",
    "print(os.listdir(PROJECT_FOLDER + \"data\"))\n",
    "print(len(os.listdir(DATA_FOLDER)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZmZ-HheL3itu"
   },
   "source": [
    "# Helper functions - seed and metric calculator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "csp2OMgo2_ar"
   },
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8V_xuoN73lON"
   },
   "source": [
    "# Dataset and transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8ucWZHeG2_as"
   },
   "outputs": [],
   "source": [
    "import albumentations as A\n",
    "import albumentations.pytorch as APT\n",
    "import cv2 \n",
    "\n",
    "train_transform = A.Compose([\n",
    "    # A.Resize(IMG_SIZE, IMG_SIZE),\n",
    "    # A.CLAHE(p=1), \n",
    "    \n",
    "    A.HorizontalFlip(p=0.75),\n",
    "    A.VerticalFlip(p=0.25),\n",
    "    A.ShiftScaleRotate(p=0.5, border_mode=cv2.BORDER_CONSTANT),\n",
    "    A.OpticalDistortion(p=0.25),\n",
    "    A.IAAPerspective(p=0.25),\n",
    "    A.CoarseDropout(p=0.5),\n",
    "\n",
    "    A.RandomBrightness(p=0.75),\n",
    "    A.ToFloat(),\n",
    "    APT.transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "\n",
    "val_transform = A.Compose([\n",
    "    # A.Resize(IMG_SIZE, IMG_SIZE),\n",
    "    # A.CLAHE(p=1),\n",
    "    A.ToFloat(),\n",
    "    APT.transforms.ToTensor(),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EiLYsfKq2_at"
   },
   "outputs": [],
   "source": [
    "class HotelTrainDataset:\n",
    "    def __init__(self, data, transform=None, data_path=\"train_images/\"):\n",
    "        self.data = data\n",
    "        self.data_path = data_path\n",
    "        self.transform = transform\n",
    "        self.fake_load = False\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        record = self.data.iloc[idx]\n",
    "        image_path = self.data_path + record[\"image\"]\n",
    "\n",
    "        if self.fake_load:\n",
    "            image = np.random.randint(0, 255, (32, 32, 3)).astype(np.uint8)\n",
    "        else:\n",
    "            image = np.array(pil_image.open(image_path)).astype(np.uint8)\n",
    "\n",
    "        if self.transform:\n",
    "            transformed = self.transform(image=image)\n",
    "        \n",
    "        return {\n",
    "            \"image\" : transformed[\"image\"],\n",
    "            \"target\" : record['hotel_id_code'],\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FpR2HfK93pvS"
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_2mse3zX3pFQ"
   },
   "outputs": [],
   "source": [
    "class EmbeddingNet(nn.Module):\n",
    "    def __init__(self, n_classes=100, embed_size=64, backbone_name=\"efficientnet_b0\"):\n",
    "        super(EmbeddingNet, self).__init__()\n",
    "\n",
    "        self.embed_size = embed_size\n",
    "        self.backbone = timm.create_model(backbone_name, pretrained=True)\n",
    "        in_features = self.backbone.get_classifier().in_features\n",
    "\n",
    "        fc_name, _ = list(self.backbone.named_modules())[-1]\n",
    "        if fc_name == 'classifier':\n",
    "            self.backbone.classifier = nn.Identity()\n",
    "        elif fc_name == 'head.fc':\n",
    "            self.backbone.head.fc = nn.Identity()\n",
    "        elif fc_name == 'fc':\n",
    "            self.backbone.fc = nn.Identity()\n",
    "        else:\n",
    "            raise Exception(\"unknown classifier layer: \" + fc_name)\n",
    "\n",
    "        self.post = nn.Sequential(\n",
    "            nn.utils.weight_norm(nn.Linear(in_features, self.embed_size*2), dim=None),\n",
    "            nn.BatchNorm1d(self.embed_size*2),\n",
    "            nn.Dropout(0.2),\n",
    "            nn.utils.weight_norm(nn.Linear(self.embed_size*2, self.embed_size)),\n",
    "        )\n",
    "\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.BatchNorm1d(self.embed_size),\n",
    "            nn.Dropout(0.2),\n",
    "            nn.Linear(self.embed_size, n_classes),\n",
    "        )\n",
    "        \n",
    "    def embed_and_classify(self, x):\n",
    "        x = self.forward(x)\n",
    "        return x, self.classifier(x)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.backbone(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.post(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mTFCinps35ci"
   },
   "source": [
    "# Model helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xW5LIe1l2_at"
   },
   "outputs": [],
   "source": [
    "def get_embeds(loader, model, bar_desc=\"Generating embeds\"):\n",
    "    targets_all = []\n",
    "    outputs_all = []\n",
    "    \n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        t = tqdm(loader, desc=bar_desc)\n",
    "        for i, sample in enumerate(t):\n",
    "            input = sample['image'].to(args.device)\n",
    "            target = sample['target'].to(args.device)\n",
    "            output = model(input)\n",
    "            \n",
    "            targets_all.extend(target.cpu().numpy())\n",
    "            outputs_all.extend(output.detach().cpu().numpy())\n",
    "            \n",
    "    return targets_all, outputs_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "syXhlJrJG2AV"
   },
   "outputs": [],
   "source": [
    "def get_distance_matrix(embeds, base_embeds, distance_func):\n",
    "    distance_matrix = []\n",
    "    base_embeds = torch.Tensor(base_embeds)\n",
    "    embeds_dataset = torch.utils.data.TensorDataset(torch.Tensor(embeds))\n",
    "    embeds_dataloader = DataLoader(embeds_dataset, num_workers=2, batch_size=1024, shuffle=False)\n",
    "    \n",
    "    t = tqdm(embeds_dataloader)\n",
    "    for i, sample in enumerate(t): \n",
    "        distances = distance_func(sample[0], base_embeds)\n",
    "        distance_matrix.extend(distances.numpy())\n",
    "        \n",
    "    return np.array(distance_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ryZ6wE0zKPiz"
   },
   "outputs": [],
   "source": [
    "def save_checkpoint(model, scheduler, optimizer, epoch, name, loss=None, score=None):\n",
    "    checkpoint = {\"epoch\": epoch,\n",
    "                  \"model\": model.state_dict(),\n",
    "                  \"scheduler\": scheduler.state_dict(),\n",
    "                  \"optimizer\": optimizer.state_dict(),\n",
    "                  \"loss\": loss,\n",
    "                  \"score\": score,\n",
    "                  }\n",
    "\n",
    "    torch.save(checkpoint, f\"{PROJECT_FOLDER}output/checkpoint-{name}.pt\")\n",
    "\n",
    "\n",
    "def load_checkpoint(model, scheduler, optimizer, name):\n",
    "    checkpoint = torch.load(f\"{PROJECT_FOLDER}output/checkpoint-{name}.pt\")\n",
    "\n",
    "    model.load_state_dict(checkpoint[\"model\"])\n",
    "    scheduler.load_state_dict(checkpoint[\"scheduler\"])\n",
    "    # do not load optimizer checkpoint, lookahead might have some gradients so it may cuz memory error\n",
    "    # optimizer.load_state_dict(checkpoint[\"optimizer\"])\n",
    "\n",
    "    return model, scheduler, optimizer, checkpoint[\"epoch\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o8sQ9dtJH1fu"
   },
   "outputs": [],
   "source": [
    "def iterate_loader(loader, epochs):\n",
    "    \"\"\"\n",
    "    Iterates through data loader with fake load (empty data) instead\n",
    "    of reading the real images to speed up. Dataloader has no state_dict\n",
    "    so manual iterating is need to get the loader with augmentations to \n",
    "    correct state.\n",
    "    \"\"\"\n",
    "    loader.dataset.fake_load = True\n",
    "    with torch.no_grad():\n",
    "        for i in range(epochs):\n",
    "            t = tqdm(loader, desc=f\"Iterating loader {i+1}/{epochs}\")\n",
    "            for j, sample in enumerate(t):\n",
    "                images = sample['image']\n",
    "                targets = sample['target']\n",
    "\n",
    "    loader.dataset.fake_load = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SntLH82s2_au"
   },
   "outputs": [],
   "source": [
    "def train_epoch(args, model, loader, criterion, optimizer, loss_optimizer, scheduler, epoch):\n",
    "    losses = []\n",
    "    targets_all = []\n",
    "    outputs_all = []\n",
    "    \n",
    "    model.train()\n",
    "    t = tqdm(loader)\n",
    "    \n",
    "    for i, sample in enumerate(t):\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        images = sample['image'].to(args.device)\n",
    "        targets = sample['target'].to(args.device)\n",
    "        \n",
    "        embeds, outputs = model.embed_and_classify(images)\n",
    "        loss = criterion(embeds, targets)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        loss_optimizer.step()\n",
    "        \n",
    "        if scheduler:\n",
    "            scheduler.step()\n",
    "                \n",
    "        losses.append(loss.item())\n",
    "        targets_all.extend(targets.cpu().numpy())\n",
    "        outputs_all.extend(torch.sigmoid(outputs).detach().cpu().numpy())\n",
    "\n",
    "        score = np.mean(targets_all == np.argmax(outputs_all, axis=1))\n",
    "        desc = f\"Epoch {epoch}/{args.epochs} - Train loss:{loss:0.4f}, score: {score:0.4f}\"\n",
    "        t.set_description(desc)\n",
    "        \n",
    "    return np.mean(losses), score\n",
    "\n",
    "\n",
    "def test_closest_match(base_df, base_embeds, valid_targets, valid_embeds, model, distance_func, closest, n_matches=5):\n",
    "    distance_matrix = get_distance_matrix(valid_embeds, base_embeds, distance_func)\n",
    "\n",
    "    preds = []\n",
    "    N_val = len(valid_embeds)\n",
    "    for i in tqdm(range(N_val), total=N_val, desc=\"Getting closest match\"):\n",
    "        tmp_df = base_df.copy()\n",
    "        tmp_df[\"distance\"] = distance_matrix[i]\n",
    "        tmp_df = tmp_df.sort_values(by=[\"distance\", \"hotel_id\"], ascending=closest).reset_index(drop=True)\n",
    "        preds.extend([tmp_df[\"hotel_id_code\"].unique()[:n_matches]])\n",
    "\n",
    "    y = np.repeat([valid_targets], repeats=n_matches, axis=0).T\n",
    "    preds = np.array(preds)\n",
    "    acc_top_1 = (preds[:, 0] == valid_targets).mean()\n",
    "    acc_top_5 = (preds == y).any(axis=1).mean()\n",
    "    print(f\"Accuracy: {acc_top_1:0.4f}, top 5 accuracy: {acc_top_5:0.4f}\")\n",
    "    return preds, distance_matrix\n",
    "\n",
    "\n",
    "def test(base_loader, valid_loader, model, distance_func, closest):\n",
    "    base_targets, base_embeds = get_embeds(base_loader, model, \"Generating embeds for train\")\n",
    "    valid_targets, valid_embeds = get_embeds(valid_loader, model, \"Generating embeds for test\")\n",
    "    val_preds, distance_matrix = test_closest_match(base_loader.dataset.data, base_embeds, valid_targets, valid_embeds, model, distance_func, closest)\n",
    "\n",
    "    return base_embeds, valid_embeds, base_targets, valid_targets, val_preds, distance_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F2xgmwBW4LjC"
   },
   "source": [
    "# Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JBkHrXYy2_av"
   },
   "outputs": [],
   "source": [
    "def sample_data(n_hotels, min_images, max_images):\n",
    "    data_df = pd.read_csv(PROJECT_FOLDER + \"data/train.csv\", parse_dates=[\"timestamp\"])\n",
    "    sample_df = data_df.groupby(\"hotel_id\").filter(lambda x: (x[\"image\"].nunique() > min_images) & (x[\"image\"].nunique() < max_images))\n",
    "    sample_df[\"hotel_id_code\"] = sample_df[\"hotel_id\"].astype('category').cat.codes.values.astype(np.int64)\n",
    "    sample_df = sample_df[sample_df[\"hotel_id_code\"] < n_hotels]\n",
    "\n",
    "    print(f\"Subsample with {len(sample_df.hotel_id.unique())} hotels out of {len(data_df.hotel_id.unique())}\" + \n",
    "          f\" with total {len(sample_df)} images ({len(sample_df) / len(data_df) * 100:0.2f} %)\")\n",
    "    \n",
    "    return sample_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Sn6HrWKQ2_aw"
   },
   "outputs": [],
   "source": [
    "# FOR TESTING DIFFERENT SETTING\n",
    "# data_df = sample_data(1000, 15, 50)\n",
    "\n",
    "# FOR FINAL TRAINING\n",
    "data_df = pd.read_csv(PROJECT_FOLDER + \"data/train.csv\", parse_dates=[\"timestamp\"])\n",
    "data_df[\"hotel_id_code\"] = data_df[\"hotel_id\"].astype('category').cat.codes.values.astype(np.int64)\n",
    "\n",
    "fig = go.Figure()\n",
    "fig.add_trace(go.Histogram(x=data_df[\"hotel_id_code\"]))\n",
    "fig.update_xaxes(type=\"category\")\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3aEmY6K7KY3H"
   },
   "outputs": [],
   "source": [
    "def train_and_validate(args, data_df):\n",
    "    model_name = f\"cosface-model-{args.backbone_name}-{IMG_SIZE}x{IMG_SIZE}-{args.embed_size}embeds-{args.n_classes}hotels\"\n",
    "    print(model_name)\n",
    "\n",
    "    seed_everything(seed=SEED)\n",
    "\n",
    "    val_df = data_df.groupby(\"hotel_id\").sample(args.val_samples, random_state=SEED)\n",
    "    train_df = data_df[~data_df[\"image\"].isin(val_df[\"image\"])]\n",
    "\n",
    "    train_dataset = HotelTrainDataset(train_df, train_transform, data_path=DATA_FOLDER)\n",
    "    train_loader = DataLoader(train_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True, drop_last=True)\n",
    "    base_dataset = HotelTrainDataset(train_df, val_transform, data_path=DATA_FOLDER)\n",
    "    base_loader = DataLoader(base_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False)\n",
    "    val_dataset = HotelTrainDataset(val_df, val_transform, data_path=DATA_FOLDER)\n",
    "    valid_loader = DataLoader(val_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False)\n",
    "\n",
    "    print(f\"Base: {len(base_dataset)}\\nValidation: {len(val_dataset)}\")\n",
    "\n",
    "    model = EmbeddingNet(args.n_classes, args.embed_size, args.backbone_name)\n",
    "    model = model.to(args.device)\n",
    "\n",
    "    distance = distances.CosineSimilarity()\n",
    "\n",
    "    criterion = losses.CosFaceLoss(num_classes=args.n_classes, embedding_size=args.embed_size, embedding_regularizer = regularizers.RegularFaceRegularizer()).to(args.device) # Accuracy: 0.7200, top 5 accuracy: 0.8460\n",
    "    loss_optimizer = torch.optim.AdamW(criterion.parameters(), lr=args.lr)\n",
    "    optimizer = Lookahead(torch.optim.AdamW(model.parameters(), lr=args.lr), k=3)\n",
    "\n",
    "    scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
    "                    optimizer,\n",
    "                    max_lr=args.lr,\n",
    "                    epochs=args.epochs,\n",
    "                    steps_per_epoch=len(train_loader),\n",
    "                    div_factor=10,\n",
    "                    final_div_factor=1,\n",
    "                    pct_start=0.1,\n",
    "                    anneal_strategy=\"cos\",\n",
    "                )\n",
    "    \n",
    "    start_epoch = 1\n",
    "\n",
    "    if args.continue_from_checkpoint:\n",
    "        model, scheduler, optimizer, last_epoch = load_checkpoint(model, scheduler, optimizer, model_name)\n",
    "        iterate_loader(train_loader, last_epoch)\n",
    "        start_epoch = start_epoch + last_epoch\n",
    "\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    for epoch in range(start_epoch, args.epochs+1):\n",
    "        train_loss, train_score = train_epoch(args, model, train_loader, criterion, optimizer, loss_optimizer, scheduler, epoch)\n",
    "        save_checkpoint(model, scheduler, optimizer, epoch, model_name, train_loss, train_score)\n",
    "        if (epoch == 1):\n",
    "            _ = test(base_loader, valid_loader, model, distance, closest=False)\n",
    "\n",
    "    base_embeds, valid_embeds, base_targets, valid_targets, val_preds, distance_matrix = test(base_loader, valid_loader, model, distance, closest=False)\n",
    "    \n",
    "    # output = {\"base_embeds\": base_embeds,\n",
    "    #           \"valid_embeds\": valid_embeds,\n",
    "    #           \"base_targets\": base_targets,\n",
    "    #           \"valid_targets\": valid_targets,\n",
    "    #           \"val_preds\": val_preds,\n",
    "    #           \"distance_matrix\": distance_matrix,\n",
    "    #           \"train_df\" : train_df,\n",
    "    #           \"valid_df\": val_df,\n",
    "    #           }\n",
    "\n",
    "    # torch.save(output, f\"{PROJECT_FOLDER}output/output-{model_name}.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EMVYKwZ64zUN"
   },
   "source": [
    "# Train and evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YONzJBtG2_a0"
   },
   "outputs": [],
   "source": [
    "# %%time \n",
    "\n",
    "# class args:\n",
    "#     epochs = 6\n",
    "#     lr = 1e-3\n",
    "#     batch_size = 32\n",
    "#     num_workers = 2\n",
    "#     embed_size = 256\n",
    "#     val_samples = 2\n",
    "#     continue_from_checkpoint = False\n",
    "#     backbone_name = \"efficientnet_b1\"\n",
    "#     n_classes = data_df[\"hotel_id_code\"].nunique()\n",
    "#     device = ('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# train_and_validate(args, data_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9VSJTgweWFxS"
   },
   "outputs": [],
   "source": [
    "# DOESNT CONVERGE\n",
    "# %%time \n",
    "\n",
    "# class args:\n",
    "#     epochs = 6\n",
    "#     lr = 1e-3\n",
    "#     batch_size = 32\n",
    "#     num_workers = 2\n",
    "#     embed_size = 256\n",
    "#     continue_from_checkpoint = False\n",
    "#     backbone_name = \"eca_nfnet_l0\"\n",
    "#     n_classes = data_df[\"hotel_id_code\"].nunique()\n",
    "#     device = ('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# train_and_validate(args, data_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "irUTFdftLHhn"
   },
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "class args:\n",
    "    epochs = 9\n",
    "    lr = 1e-3\n",
    "    batch_size = 24\n",
    "    num_workers = 2\n",
    "    embed_size = 4096\n",
    "    val_samples = 1\n",
    "    continue_from_checkpoint = True\n",
    "    backbone_name = \"ecaresnet50d_pruned\"\n",
    "    n_classes = data_df[\"hotel_id_code\"].nunique()\n",
    "    device = ('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "train_and_validate(args, data_df)\n",
    "\n",
    "# RESULTS\n",
    "# epoch 1: Accuracy: 0.3934, top 5 accuracy: 0.5430\n",
    "# Epoch 1/9 - Train loss:27.5661, score: 0.0001: 100%|██████████| 3741/3741 [1:59:16<00:00,  1.91s/it]\n",
    "# Iterating loader 1/2: 100%|██████████| 3741/3741 [18:17<00:00,  3.41it/s]\n",
    "# Iterating loader 2/2: 100%|██████████| 3741/3741 [17:48<00:00,  3.50it/s]\n",
    "# Epoch 3/9 - Train loss:27.5679, score: 0.0002: 100%|██████████| 3741/3741 [2:01:29<00:00,  1.95s/it]\n",
    "# Epoch 4/9 - Train loss:22.9861, score: 0.0001:  96%|█████████▌| 3588/3741 [1:54:30<07:12,  2.83s/it]\n",
    "# Iterating loader 1/3: 100%|██████████| 3741/3741 [01:17<00:00, 48.49it/s]\n",
    "# ...\n",
    "# Iterating loader 3/3: 100%|██████████| 3741/3741 [01:17<00:00, 48.10it/s]\n",
    "# Epoch 4/9 - Train loss:26.0165, score: 0.0001: 100%|██████████| 3741/3741 [1:59:41<00:00,  1.92s/it]\n",
    "# Epoch 5/9 - Train loss:19.6462, score: 0.0001:  80%|███████▉  | 2981/3741 [1:28:00<32:16,  2.55s/it]\n",
    "# Iterating loader 1/4: 100%|██████████| 3741/3741 [01:35<00:00, 39.11it/s]\n",
    "# ...\n",
    "# Iterating loader 4/4: 100%|██████████| 3741/3741 [01:42<00:00, 36.44it/s]\n",
    "# Epoch 5/9 - Train loss:23.3521, score: 0.0002: 100%|██████████| 3741/3741 [1:28:45<00:00,  1.42s/it]\n",
    "# Epoch 6/9 - Train loss:15.9608, score: 0.0002: 100%|██████████| 3741/3741 [1:28:14<00:00,  1.42s/it]\n",
    "# Epoch 7/9 - Train loss:14.7336, score: 0.0002: 100%|██████████| 3741/3741 [1:30:35<00:00,  1.45s/it]\n",
    "# Epoch 8/9 - Train loss:13.4082, score: 0.0000:   3%|▎         | 101/3741 [01:13<40:14,  1.51it/s]\n",
    "# Iterating loader 1/7: 100%|██████████| 3741/3741 [01:18<00:00, 47.94it/s]\n",
    "# ...\n",
    "# Iterating loader 7/7: 100%|██████████| 3741/3741 [01:15<00:00, 49.27it/s]\n",
    "# Epoch 8/9 - Train loss:17.9616, score: 0.0001: 100%|██████████| 3741/3741 [1:59:38<00:00,  1.92s/it]\n",
    "# Epoch 9/9 - Train loss:10.6147, score: 0.0001: 100%|██████████| 3741/3741 [2:01:29<00:00,  1.95s/it]\n",
    "# Generating embeds for train: 100%|██████████| 3742/3742 [18:36<00:00,  3.35it/s]\n",
    "# Generating embeds for test: 100%|██████████| 324/324 [01:41<00:00,  3.19it/s]\n",
    "# 100%|██████████| 8/8 [01:23<00:00, 10.50s/it]\n",
    "# Getting closest match: 100%|██████████| 7770/7770 [08:01<00:00, 16.14it/s]\n",
    "# Accuracy: 0.6542, top 5 accuracy: 0.7889"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "hotel-id-cosface-training.ipynb",
   "provenance": []
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}