{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master')\nimport timm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-06-03T14:15:53.312128Z","iopub.execute_input":"2022-06-03T14:15:53.312384Z","iopub.status.idle":"2022-06-03T14:15:53.316318Z","shell.execute_reply.started":"2022-06-03T14:15:53.312356Z","shell.execute_reply":"2022-06-03T14:15:53.315633Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Shared Imports\nimport random\nimport os\nimport pathlib\nfrom typing import Iterator, List, Optional, Tuple\nimport json\n\nimport math\nimport numpy as np\nimport pandas as pd\nimport matplotlib\nimport matplotlib.pyplot as plt\nfrom PIL import Image as pil_image\n# model imports\n## Pytorch/modell stuff\nimport torch\nimport torch.nn as nn\nfrom torchmetrics import Accuracy\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor\nfrom pytorch_lightning.loggers import TensorBoardLogger\n\nfrom sklearn.metrics.pairwise import cosine_similarity\nfrom sklearn.model_selection import train_test_split\n\n# Pre-processing\nimport albumentations as A\nimport albumentations.pytorch as APT\nimport cv2 \nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.326946Z","iopub.execute_input":"2022-06-03T14:15:53.327432Z","iopub.status.idle":"2022-06-03T14:15:53.334229Z","shell.execute_reply.started":"2022-06-03T14:15:53.327403Z","shell.execute_reply":"2022-06-03T14:15:53.333568Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Globals","metadata":{}},{"cell_type":"code","source":"SEED = 42\nNUM_WORKERS = 2\n\n# Wheter to PAD the images\nPAD = True\n# The size of the images\nPATCH = (256, 256)\n# The number of matches to consider\nN_MATCHES = 5\n# From trainning to ensure the same val-train division\nVAL_SIZE = 0.1\nBATCH_SIZE = 16\n# The used embedding size\nEMBEDDING_SIZE = 4096\n# The used base model\nBASE_MODEL = \"efficientnet_b1\"\n# Set random Seed\npl.seed_everything(SEED)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.342994Z","iopub.execute_input":"2022-06-03T14:15:53.343781Z","iopub.status.idle":"2022-06-03T14:15:53.354245Z","shell.execute_reply.started":"2022-06-03T14:15:53.343742Z","shell.execute_reply":"2022-06-03T14:15:53.353428Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# directory of the data\nDATA_DIR = pathlib.Path(\"../input/mlip-pad-resize-256x256\")\n# Work directory, where to store the data\nWORKING_DIR = pathlib.Path(\"\")\n# Locations of the original train set, to derive the chain names\nCHAIN_DIR = pathlib.Path(\"../input/hotel-id-to-combat-human-trafficking-2022-fgvc9/train_images\")\n# Locations of the train images in the data directory\nTRAIN_DIR = DATA_DIR / pathlib.Path(\"pad_and_resize/train_images\")\n# Directory of the model weights and saved embeddings\nMODEL_WEIGHTS_DIR = pathlib.Path(\"../input/mlip-hotelid-sim-weights-emb\")\nMODEL_WEIGHTS = pathlib.Path(\"../input/mlip-hotelid-sim-weights-emb/logs/lightning_logs/version_0/checkpoints/epoch_0009.step_000025139.val-map_6.4391.last.ckpt\")\nBASE_EMB = MODEL_WEIGHTS_DIR / \"base_image-embeddings.pkl\"","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.368421Z","iopub.execute_input":"2022-06-03T14:15:53.369053Z","iopub.status.idle":"2022-06-03T14:15:53.374370Z","shell.execute_reply.started":"2022-06-03T14:15:53.369017Z","shell.execute_reply":"2022-06-03T14:15:53.373712Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loading data","metadata":{}},{"cell_type":"code","source":"hotel_id_code_df = pd.read_csv(MODEL_WEIGHTS_DIR / 'hotel_id_code_mapping.csv')\nhotel_id_code_map = hotel_id_code_df.set_index('hotel_code').to_dict()[\"hotel_id\"]\nhotel_id_code_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.380194Z","iopub.execute_input":"2022-06-03T14:15:53.380862Z","iopub.status.idle":"2022-06-03T14:15:53.398483Z","shell.execute_reply.started":"2022-06-03T14:15:53.380827Z","shell.execute_reply":"2022-06-03T14:15:53.397816Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"chain_names = os.listdir(CHAIN_DIR)\ntrain_file = DATA_DIR / pathlib.Path(\"train.csv\")\n\n# Encode the chain identifiers so that the model can work with it and save it so it can be retrieved\ntrain_df = pd.read_csv(train_file)    \ntrain_df[\"hotel_code\"] = train_df[\"hotel_id\"].astype('category').cat.codes.values.astype(np.int64)\n\nhotel_id_code_df = train_df.drop(columns=[\"image_id\"]).drop_duplicates().reset_index(drop=True)\nhotel_id_code_df.to_csv(WORKING_DIR / 'hotel_id_code_mapping.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.399830Z","iopub.execute_input":"2022-06-03T14:15:53.400208Z","iopub.status.idle":"2022-06-03T14:15:53.445530Z","shell.execute_reply.started":"2022-06-03T14:15:53.400176Z","shell.execute_reply":"2022-06-03T14:15:53.444917Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"Number of images:\", len(train_df))\nprint(\"Number of different classes:\", len(chain_names))\ntrain_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.446978Z","iopub.execute_input":"2022-06-03T14:15:53.447339Z","iopub.status.idle":"2022-06-03T14:15:53.458241Z","shell.execute_reply.started":"2022-06-03T14:15:53.447306Z","shell.execute_reply":"2022-06-03T14:15:53.457454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Split the validation and train set\ntrain_set, val_set = train_test_split(train_df, test_size=VAL_SIZE, random_state=SEED)\nprint(\"Number of train images:\", len(train_set))\nprint(\"Number of val images:\", len(val_set))\ntrain_set.head()","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.459401Z","iopub.execute_input":"2022-06-03T14:15:53.459801Z","iopub.status.idle":"2022-06-03T14:15:53.476280Z","shell.execute_reply.started":"2022-06-03T14:15:53.459768Z","shell.execute_reply":"2022-06-03T14:15:53.475501Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Pre-process","metadata":{}},{"cell_type":"code","source":"def open_preprocess_img(img_path: pathlib.Path):\n    img = cv2.imread(str(img_path))\n    \n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    if PAD: img = pad(img)\n    \n    return cv2.resize(img, PATCH)\n\ndef pad(img):\n    w, h, c = np.shape(img)\n    const = 0\n        \n    if w == h: return img\n    elif (w - h) % 2 != 0: const = 1\n        \n    if w < h:\n        half_py = (h - w) // 2       \n        return cv2.copyMakeBorder(img, 0, 0, half_py, half_py + const, cv2.BORDER_CONSTANT, value=0)\n    elif h < w:\n        half_px = (w - h) // 2\n        return cv2.copyMakeBorder(img, half_px, half_px + const, 0, 0, cv2.BORDER_CONSTANT, value=0)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.477772Z","iopub.execute_input":"2022-06-03T14:15:53.478072Z","iopub.status.idle":"2022-06-03T14:15:53.487081Z","shell.execute_reply.started":"2022-06-03T14:15:53.478037Z","shell.execute_reply":"2022-06-03T14:15:53.484977Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Augmentations","metadata":{}},{"cell_type":"code","source":"base_transform = A.Compose([\n    A.ToFloat(),\n    APT.transforms.ToTensorV2(),\n])\n\nval_aug = A.Compose([\n    A.CoarseDropout(p=0.75, max_holes=1, \n                    min_height=PATCH[0]//4, max_height=PATCH[0]//2,\n                    min_width=PATCH[1]//4,  max_width=PATCH[1]//2, \n                    fill_value=(255,0,0)),# simulating occlusions\n    A.ToFloat(),\n    APT.transforms.ToTensorV2(),\n])","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.488817Z","iopub.execute_input":"2022-06-03T14:15:53.489632Z","iopub.status.idle":"2022-06-03T14:15:53.496579Z","shell.execute_reply.started":"2022-06-03T14:15:53.489597Z","shell.execute_reply":"2022-06-03T14:15:53.495734Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataloader","metadata":{}},{"cell_type":"code","source":"class ImageDataset(Dataset):\n    def __init__(self,\n                 data: pd.DataFrame,\n                 data_path: pathlib.Path,\n                 transform: Optional = None,\n                ):\n        self.data = data\n        self.data_path = data_path\n        self.transform = transform\n\n    def __len__(self) -> int:\n        return len(self.data)\n\n    def __getitem__(self, idx: int):\n        record = self.data.iloc[idx]\n\n        image_path = self.data_path / record[\"image_id\"]\n        image = np.array(pil_image.open(image_path)).astype(np.uint8)\n        \n        if self.transform:\n            transformed = self.transform(image=image)\n            image = transformed[\"image\"]\n        \n        label = record['hotel_code']\n        return image, label","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.498086Z","iopub.execute_input":"2022-06-03T14:15:53.498857Z","iopub.status.idle":"2022-06-03T14:15:53.508766Z","shell.execute_reply.started":"2022-06-03T14:15:53.498814Z","shell.execute_reply":"2022-06-03T14:15:53.507984Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"markdown","source":"## Arcface models","metadata":{}},{"cell_type":"code","source":"# source: https://github.com/ronghuaiyang/arcface-pytorch/blob/master/models/metrics.py\nclass 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):\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.weight = nn.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)).clamp(0, 1))\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        # -------------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\n        return output","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.510495Z","iopub.execute_input":"2022-06-03T14:15:53.511098Z","iopub.status.idle":"2022-06-03T14:15:53.523855Z","shell.execute_reply.started":"2022-06-03T14:15:53.511063Z","shell.execute_reply":"2022-06-03T14:15:53.523086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HotelModelArcface(pl.LightningModule):\n    def __init__(self,\n                n_hotels: int,\n                steps_per_epoch: int,\n                n_embeddings: int = 256,\n                base_model = None,\n                pretrained: bool = False,\n                learning_rate: float = 0.003,\n                \n                ):\n        super().__init__()\n        \n        # Hyperparams\n        self.n_embeddings = n_embeddings\n        self.n_hotels = n_hotels\n        self.learning_rate = learning_rate\n        self.steps_per_epoch = steps_per_epoch\n        \n        # Metrics\n        self.loss_fn = nn.CrossEntropyLoss()\n        self.train_acc = Accuracy()\n        self.val_acc = Accuracy()\n        \n        # Model Definition \n        ## Base model\n        self.base_model = timm.create_model(base_model, pretrained=False)        \n        in_features = self.base_model.get_classifier().in_features\n        \n        fc_name, _ = list(self.base_model.named_modules())[-1]\n        if fc_name == 'classifier':\n            self.base_model.classifier = nn.Identity()\n        elif fc_name == 'head.fc':\n            self.base_model.head.fc = nn.Identity()\n        elif fc_name == 'fc':\n            self.base_model.fc = nn.Identity()\n        else:\n            raise Exception(\"unknown classifier layer: \" + fc_name)\n        \n        ## Arcface module\n        # self.arc_face = ArcMarginProduct(self.n_embeddings, n_hotels, s=30.0, m=0.20, easy_margin=False)\n        \n        ## Top model\n        self.top_model = nn.Sequential(\n            nn.utils.weight_norm(nn.Linear(in_features, self.n_embeddings*2), dim=None),\n            nn.BatchNorm1d(self.n_embeddings*2),\n            nn.Dropout(0.2),\n            nn.utils.weight_norm(nn.Linear(self.n_embeddings*2, self.n_embeddings)),\n            nn.BatchNorm1d(self.n_embeddings),\n        )\n        \n        # Save hyper params\n        self.save_hyperparameters()\n\n    def configure_optimizers(self):\n        optimizer = Lookahead(torch.optim.AdamW(self.parameters(), lr=self.learning_rate), k=3)\n        \n        scheduler = torch.optim.lr_scheduler.OneCycleLR(\n                    optimizer,\n                    max_lr=self.learning_rate,\n                    epochs=EPOCHS,\n                    steps_per_epoch=self.steps_per_epoch,\n                    div_factor=10,\n                    final_div_factor=1,\n                    pct_start=0.1,\n                    anneal_strategy=\"cos\",\n                )\n        \n        schedule = {\n            # Required: the scheduler instance.\n            \"scheduler\": scheduler,\n        }\n        return [optimizer], [schedule]\n    \n    def forward(self, x, targets = None):\n        y_hat = self.base_model(x)\n        y_hat = y_hat.view(y_hat.size(0), -1)\n        y_hat = self.top_model(y_hat)\n        \n        if targets is not None:\n            y_hat = self.arc_face(y_hat, targets)\n\n        return y_hat\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        \n        # Forward pass\n        y_hat = self.forward(x, y)\n        loss = self.loss_fn(y_hat, y)\n        self.train_acc(y_hat, y)\n\n        # Store results\n        self.log(\"train_loss\", loss, prog_bar=False)\n        \n        return loss\n    \n    def training_epoch_end(self, train_step_outputs) -> None:\n        # Log metrics\n        self.log(\"train_acc\", self.train_acc, prog_bar=True)\n\n    def validation_step(self, batch, batch_idx):\n        x, y = batch\n        \n        # Forward pass\n        y_hat = self.forward(x, y)\n        loss = self.loss_fn(y_hat, y)\n        self.val_acc(y_hat, y)\n\n        # Store results\n        self.log(\"val_loss\", loss, prog_bar=False)\n        return y_hat\n        \n    def validation_epoch_end(self, validation_step_outputs) -> None:\n        self.log(\"val_acc\", self.val_acc, prog_bar=True)\n        \n    def predict_step(self, batch, batch_idx):\n        y_hat = self.forward(batch)\n        return y_hat\n    \n    def test_step(self, batch, batch_idx):\n        # Forward pass\n        x, y = batch\n        sample_emb = self.forward(x)\n        \n        # Cosine sim\n        cosine_sim = F.cosine_similarity(sample_emb, self.base_emb)\n        sorted_idx = torch.argsort(cosine_sim, descending=True)\n        sorted_lbls = self.base_labels[sorted_idx]\n\n        # Because unique does not return properly\n        top_lbls = []\n        i = 0\n        while len(top_lbls) < 5:\n            if sorted_lbls[i] not in top_lbls:\n                top_lbls.append(sorted_lbls[i])\n            i += 1\n        top_lbls = torch.Tensor(top_lbls).long().to(('cuda' if torch.cuda.is_available() else 'cpu'))\n        \n        # Calculate MAP@5\n        y_idx = (y == top_lbls).nonzero(as_tuple=True)[0]\n        if len(y_idx) == 0:\n            map5 = torch.zeros(1).to(('cuda' if torch.cuda.is_available() else 'cpu'))\n        else:\n            map5 = 1 / (y_idx+1)\n\n        acc_top_1 = (y == sorted_lbls[0])\n        acc_top_5 = (y == top_lbls[:N_MATCHES])\n        return map5, acc_top_1.float(), acc_top_5.float()\n        \n    def test_epoch_end(self, test_step_outputs) -> None:\n        map5 = torch.stack([ap for ap, _, _ in test_step_outputs])\n        acc_top_1 = torch.stack([acc1 for _, acc1, _ in test_step_outputs])\n        acc_top_5 = torch.stack([acc5 for _, _, acc5 in test_step_outputs])\n        print(len(map5), len(acc_top_1), len(acc_top_5))\n        print(torch.mean(map5))\n        print(torch.mean(acc_top_1))\n        print(torch.mean(acc_top_5.any(axis=1).float()))\n        \n        self.log(\"train_map5\", torch.mean(map5), prog_bar=True)\n        self.log(\"train_acc\", torch.mean(acc_top_1), prog_bar=True)\n        self.log(\"train-acc5\", torch.mean(acc_top_5.any(axis=1).float()), prog_bar=True)\n        \n    def set_base_emb(self, base_emb, base_labels):\n        self.base_emb = base_emb\n        self.base_labels = base_labels","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.525484Z","iopub.execute_input":"2022-06-03T14:15:53.526218Z","iopub.status.idle":"2022-06-03T14:15:53.556840Z","shell.execute_reply.started":"2022-06-03T14:15:53.526185Z","shell.execute_reply":"2022-06-03T14:15:53.556167Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Base model","metadata":{}},{"cell_type":"code","source":"class HotelModel(pl.LightningModule):\n    def __init__(self,\n                n_hotels: int,\n                steps_per_epoch: int,\n                n_embeddings: int = 256,\n                base_model = None,\n                pretrained: bool = False,\n                learning_rate: float = 0.003,\n                \n                ):\n        super().__init__()\n        \n        # Hyperparams\n        self.n_embeddings = n_embeddings\n        self.n_hotels = n_hotels\n        self.learning_rate = learning_rate\n        self.steps_per_epoch = steps_per_epoch\n        \n        # Metrics\n        # self.loss_fn = F.cross_entropy\n        self.train_acc = Accuracy()\n        self.val_acc = Accuracy()\n        \n        # Model Definition \n        ## Base model\n        self.base_model = timm.create_model(base_model, pretrained=False)        \n        in_features = self.base_model.get_classifier().in_features\n        \n        fc_name, _ = list(self.base_model.named_modules())[-1]\n        if fc_name == 'classifier':\n            self.base_model.classifier = nn.Identity()\n        elif fc_name == 'head.fc':\n            self.base_model.head.fc = nn.Identity()\n        elif fc_name == 'fc':\n            self.base_model.fc = nn.Identity()\n        else:\n            raise Exception(\"unknown classifier layer: \" + fc_name)\n        \n        ## Top model\n        self.top_model = nn.Sequential(\n            nn.utils.weight_norm(nn.Linear(in_features, self.n_embeddings*2), dim=None),\n            nn.BatchNorm1d(self.n_embeddings*2),\n            nn.Dropout(0.2),\n            nn.utils.weight_norm(nn.Linear(self.n_embeddings*2, self.n_embeddings)),\n        )\n        \n        ## Classifier\n        self.classifier = nn.Sequential(\n            nn.BatchNorm1d(self.n_embeddings),\n            nn.Dropout(0.2),\n            nn.Linear(self.n_embeddings, self.n_hotels),\n        )\n        \n        # Save hyper params\n        self.save_hyperparameters()\n\n    def configure_optimizers(self):\n        optimizer = Lookahead(torch.optim.AdamW(self.parameters(), lr=self.learning_rate), k=3)\n        \n        scheduler = torch.optim.lr_scheduler.OneCycleLR(\n                    optimizer,\n                    max_lr=self.learning_rate,\n                    epochs=EPOCHS,\n                    steps_per_epoch=self.steps_per_epoch,\n                    div_factor=10,\n                    final_div_factor=1,\n                    pct_start=0.1,\n                    anneal_strategy=\"cos\",\n                )\n        \n        schedule = {\n            # Required: the scheduler instance.\n            \"scheduler\": scheduler,\n        }\n        return [optimizer], [schedule]\n    \n    def forward(self, x):\n        emb = self.base_model(x)\n        emb = emb.view(emb.size(0), -1)\n        emb = self.top_model(emb)\n        y_hat = self.classifier(emb)\n        return emb\n\n    def test_step(self, batch, batch_idx):\n        # Forward pass\n        x, y = batch\n        sample_emb = self.forward(x)\n        \n        # Cosine sim\n        cosine_sim = F.cosine_similarity(sample_emb, self.base_emb)\n        sorted_idx = torch.argsort(cosine_sim, descending=True)\n        sorted_lbls = self.base_labels[sorted_idx]\n\n        # Because unique does not return properly\n        top_lbls = []\n        i = 0\n        while len(top_lbls) < 5:\n            if sorted_lbls[i] not in top_lbls:\n                top_lbls.append(sorted_lbls[i])\n            i += 1\n        top_lbls = torch.Tensor(top_lbls).long().to(('cuda' if torch.cuda.is_available() else 'cpu'))\n        \n        # Calculate MAP@5\n        y_idx = (y == top_lbls).nonzero(as_tuple=True)[0]\n        if len(y_idx) == 0:\n            map5 = torch.zeros(1).to(('cuda' if torch.cuda.is_available() else 'cpu'))\n        else:\n            map5 = 1 / (y_idx+1)\n\n        acc_top_1 = (y == sorted_lbls[0])\n        acc_top_5 = (y == top_lbls[:N_MATCHES])\n        return map5, acc_top_1.float(), acc_top_5.float()\n        \n    def test_epoch_end(self, test_step_outputs) -> None:\n        map5 = torch.stack([ap for ap, _, _ in test_step_outputs])\n        acc_top_1 = torch.stack([acc1 for _, acc1, _ in test_step_outputs])\n        acc_top_5 = torch.stack([acc5 for _, _, acc5 in test_step_outputs])\n        print(len(map5), len(acc_top_1), len(acc_top_5))\n        print(torch.mean(map5))\n        print(torch.mean(acc_top_1))\n        print(torch.mean(acc_top_5.any(axis=1).float()))\n        \n        self.log(\"train_map5\", torch.mean(map5), prog_bar=True)\n        self.log(\"train_acc\", torch.mean(acc_top_1), prog_bar=True)\n        self.log(\"train-acc5\", torch.mean(acc_top_5.any(axis=1).float()), prog_bar=True)\n        \n    def set_base_emb(self, base_emb, base_labels):\n        self.base_emb = base_emb\n        self.base_labels = base_labels","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.558327Z","iopub.execute_input":"2022-06-03T14:15:53.558682Z","iopub.status.idle":"2022-06-03T14:15:53.583732Z","shell.execute_reply.started":"2022-06-03T14:15:53.558651Z","shell.execute_reply":"2022-06-03T14:15:53.582738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loaders for the model","metadata":{}},{"cell_type":"code","source":"def get_arc_model(model_type, backbone_name, checkpoint_path, args):\n    model = HotelModelArcface.load_from_checkpoint(checkpoint_path, map_location='cpu', strict=False)\n    return model","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.585307Z","iopub.execute_input":"2022-06-03T14:15:53.586039Z","iopub.status.idle":"2022-06-03T14:15:53.595221Z","shell.execute_reply.started":"2022-06-03T14:15:53.585961Z","shell.execute_reply":"2022-06-03T14:15:53.594500Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model(model_type, backbone_name, checkpoint_path, args):\n    model = HotelModel.load_from_checkpoint(checkpoint_path, map_location='cpu', strict=False)\n    return model","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.596475Z","iopub.execute_input":"2022-06-03T14:15:53.597189Z","iopub.status.idle":"2022-06-03T14:15:53.603737Z","shell.execute_reply.started":"2022-06-03T14:15:53.597094Z","shell.execute_reply":"2022-06-03T14:15:53.603000Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Helper Functions","metadata":{}},{"cell_type":"code","source":"def generate_torch_embeddings(loader, model, bar_desc=\"Generating embeds\"):\n    targets_all = []\n    outputs_all = []\n    \n    model = model.to(args.device)\n    model.eval()\n    with torch.no_grad():\n        x = tqdm(loader, desc=bar_desc)\n        for i, sample in enumerate(x):\n            input = sample[0].to(args.device)\n            target = sample[1].to(args.device)\n            output = model(input)\n            \n            targets_all.extend(target.cpu().numpy())\n            outputs_all.extend(output.detach().cpu())\n\n    targets_all = np.array(targets_all)\n    print(len(outputs_all))\n    outputs_all = torch.stack(outputs_all)\n    print(outputs_all.shape)\n            \n    return outputs_all, targets_all","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.605693Z","iopub.execute_input":"2022-06-03T14:15:53.606106Z","iopub.status.idle":"2022-06-03T14:15:53.614314Z","shell.execute_reply.started":"2022-06-03T14:15:53.606072Z","shell.execute_reply":"2022-06-03T14:15:53.613640Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generate the embeddings and test the model","metadata":{}},{"cell_type":"code","source":"class args:\n    epochs = 5\n    lr = 1e-3\n    batch_size = 64\n    num_workers = 2\n    val_samples = 1\n    embedding_size = 128\n    backbone_name = \"efficientnet_b0\"\n    n_classes = train_df[\"hotel_code\"].nunique()\n    device = ('cuda' if torch.cuda.is_available() else 'cpu')\n    \npattern = \"epoch_{epoch:04d}.step_{step:09d}.val-map_{val_loss:.4f}\"\nModelCheckpoint.CHECKPOINT_NAME_LAST = pattern + \".last\"\ncheckpointer = ModelCheckpoint(\n        monitor=\"val_acc\",\n        filename=pattern + \".best\",\n        save_last=True,\n        auto_insert_metric_name=False,\n        save_top_k=1,\n    )\n\ntrainer = pl.Trainer(\n    max_epochs=1,\n    gpus=torch.cuda.device_count(),\n    callbacks=[checkpointer, LearningRateMonitor()],\n    default_root_dir=\"logs/\",\n)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.615447Z","iopub.execute_input":"2022-06-03T14:15:53.615918Z","iopub.status.idle":"2022-06-03T14:15:53.627844Z","shell.execute_reply.started":"2022-06-03T14:15:53.615883Z","shell.execute_reply":"2022-06-03T14:15:53.626978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_data = ImageDataset(train_set, TRAIN_DIR, transform=base_transform)\nval_data = ImageDataset(val_set, TRAIN_DIR, transform=val_aug)\n\ntrain_dataloader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, drop_last=True)\nval_dataloader = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=NUM_WORKERS)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.629729Z","iopub.execute_input":"2022-06-03T14:15:53.630022Z","iopub.status.idle":"2022-06-03T14:15:53.639175Z","shell.execute_reply.started":"2022-06-03T14:15:53.629988Z","shell.execute_reply":"2022-06-03T14:15:53.638377Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = get_arc_model(\"classification\", \n                  BASE_MODEL,\n                  MODEL_WEIGHTS, \n                  args)\n\nbase_embeds, base_targets = generate_torch_embeddings(train_dataloader, model, \"Generate base embeddings\")\n\nmodel.set_base_emb(base_embeds.to(\"cuda\"), torch.from_numpy(base_targets).to(\"cuda\"))\n\ntrainer.test(model, dataloaders=val_dataloader)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:15:53.640462Z","iopub.execute_input":"2022-06-03T14:15:53.640874Z","iopub.status.idle":"2022-06-03T14:21:36.359272Z","shell.execute_reply.started":"2022-06-03T14:15:53.640831Z","shell.execute_reply":"2022-06-03T14:21:36.358597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model specific","metadata":{}},{"cell_type":"markdown","source":"## Arcface","metadata":{}},{"cell_type":"code","source":"MODEL_WEIGHTS = pathlib.Path(\"../input/mlip-hotelid-sim-weights-emb/logs/lightning_logs/version_0/checkpoints/epoch_0009.step_000025139.val-map_6.4391.last.ckpt\")\n\nmodel = get_arc_model(\"classification\", \n                  BASE_MODEL,\n                  MODEL_WEIGHTS, \n                  args)\n\nbase_embeds, base_targets = generate_torch_embeddings(train_dataloader, model, \"Generate base embeddings\")\n\nmodel.set_base_emb(base_embeds.to(\"cuda\"), torch.from_numpy(base_targets).to(\"cuda\"))\n\ntrainer.test(model, dataloaders=val_dataloader)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:21:36.361413Z","iopub.execute_input":"2022-06-03T14:21:36.361677Z","iopub.status.idle":"2022-06-03T14:23:45.956876Z","shell.execute_reply.started":"2022-06-03T14:21:36.361642Z","shell.execute_reply":"2022-06-03T14:23:45.954724Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cosface","metadata":{}},{"cell_type":"code","source":"MODEL_WEIGHTS = pathlib.Path(\"../input/cosfaceweights/logs/lightning_logs/version_0/checkpoints/epoch_0009.step_000025139.val-acc_0.0002.last.ckpt\")\n\nmodel = get_model(\"classification\", \n                  BASE_MODEL,\n                  MODEL_WEIGHTS, \n                  args)\n\nbase_embeds, base_targets = generate_torch_embeddings(train_dataloader, model, \"Generate base embeddings\")\n\nmodel.set_base_emb(base_embeds.to(\"cuda\"), torch.from_numpy(base_targets).to(\"cuda\"))\n\ntrainer.test(model, dataloaders=val_dataloader)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:23:45.958173Z","iopub.status.idle":"2022-06-03T14:23:45.958498Z","shell.execute_reply.started":"2022-06-03T14:23:45.958343Z","shell.execute_reply":"2022-06-03T14:23:45.958361Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Tripletloss","metadata":{}},{"cell_type":"code","source":"MODEL_WEIGHTS = pathlib.Path(\"../input/tripletloss-weights/logs/lightning_logs/version_0/checkpoints/epoch_0009.step_000025139.val-map_0.1461.last.ckpt\")\n\nmodel = get_model(\"classification\", \n                  BASE_MODEL,\n                  MODEL_WEIGHTS, \n                  args)\n\nbase_embeds, base_targets = generate_torch_embeddings(train_dataloader, model, \"Generate base embeddings\")\n\nmodel.set_base_emb(base_embeds.to(\"cuda\"), torch.from_numpy(base_targets).to(\"cuda\"))\n\ntrainer.test(model, dataloaders=val_dataloader)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:23:45.959968Z","iopub.status.idle":"2022-06-03T14:23:45.960497Z","shell.execute_reply.started":"2022-06-03T14:23:45.960263Z","shell.execute_reply":"2022-06-03T14:23:45.960286Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Crossent","metadata":{}},{"cell_type":"code","source":"MODEL_WEIGHTS = pathlib.Path(\"../input/crossent-weights/logs/lightning_logs/version_0/checkpoints/epoch_0009.step_000025139.val-map_2.7154.last.ckpt\")\n\nmodel = get_model(\"classification\", \n                  BASE_MODEL,\n                  MODEL_WEIGHTS, \n                  args)\n\nbase_embeds, base_targets = generate_torch_embeddings(train_dataloader, model, \"Generate base embeddings\")\n\nmodel.set_base_emb(base_embeds.to(\"cuda\"), torch.from_numpy(base_targets).to(\"cuda\"))\n\ntrainer.test(model, dataloaders=val_dataloader)","metadata":{"execution":{"iopub.status.busy":"2022-06-03T14:23:45.961980Z","iopub.status.idle":"2022-06-03T14:23:45.962901Z","shell.execute_reply.started":"2022-06-03T14:23:45.962662Z","shell.execute_reply":"2022-06-03T14:23:45.962685Z"},"trusted":true},"execution_count":null,"outputs":[]}]}