{"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":"# Packages and Imports","metadata":{}},{"cell_type":"code","source":"!git clone https://github.com/rwightman/pytorch-image-models\n!cd pytorch-image-models && pip install -e .\n\nimport sys\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master')","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:02.864827Z","iopub.execute_input":"2022-06-06T06:59:02.865590Z","iopub.status.idle":"2022-06-06T06:59:17.779042Z","shell.execute_reply.started":"2022-06-06T06:59:02.865499Z","shell.execute_reply":"2022-06-06T06:59:17.778190Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np \nimport pandas as pd\nimport random\nimport os\nimport math\nimport time\n\nfrom sklearn.metrics.pairwise import cosine_similarity\nfrom sklearn.utils.class_weight import compute_class_weight\nfrom sklearn.model_selection import train_test_split\nfrom PIL import Image as pil_image\nfrom tqdm import tqdm\nimport scipy\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\n\nimport timm","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:17.780848Z","iopub.execute_input":"2022-06-06T06:59:17.781146Z","iopub.status.idle":"2022-06-06T06:59:22.279089Z","shell.execute_reply.started":"2022-06-06T06:59:17.781120Z","shell.execute_reply":"2022-06-06T06:59:22.278138Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Global","metadata":{}},{"cell_type":"code","source":"IMG_SIZE = 512\nSEED = 42\nN_MATCHES = 5\nUSE_TPU = False\n\nPROJECT_FOLDER = \"../input/hotel-id-to-combat-human-trafficking-2022-fgvc9/\"\nDATA_FOLDER = f\"../input/hotelid-2022-train-images-{IMG_SIZE}x{IMG_SIZE}/\"\nIMAGE_FOLDER = DATA_FOLDER+\"images/\"\nOUTPUT_FOLDER = \"\"\n\n# for TPU\nif USE_TPU:\n    !pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.7-cp37-cp37m-linux_x86_64.whl\n    import torch_xla\n    import torch_xla.core.xla_model as xm\ntorch.set_default_tensor_type('torch.FloatTensor')","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:22.280932Z","iopub.execute_input":"2022-06-06T06:59:22.281241Z","iopub.status.idle":"2022-06-06T06:59:22.289353Z","shell.execute_reply.started":"2022-06-06T06:59:22.281211Z","shell.execute_reply":"2022-06-06T06:59:22.288037Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(os.listdir(PROJECT_FOLDER))\nprint(len(os.listdir(DATA_FOLDER)))","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:22.292054Z","iopub.execute_input":"2022-06-06T06:59:22.292791Z","iopub.status.idle":"2022-06-06T06:59:22.305660Z","shell.execute_reply.started":"2022-06-06T06:59:22.292750Z","shell.execute_reply":"2022-06-06T06:59:22.304652Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_df = pd.read_csv(DATA_FOLDER+\"train.csv\")\n# encode hotel ids\ndata_df[\"hotel_id_code\"] = data_df[\"hotel_id\"].astype('category').cat.codes.values.astype(np.int64)","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:22.307456Z","iopub.execute_input":"2022-06-06T06:59:22.307772Z","iopub.status.idle":"2022-06-06T06:59:22.367198Z","shell.execute_reply.started":"2022-06-06T06:59:22.307737Z","shell.execute_reply":"2022-06-06T06:59:22.366436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class args:\n    epochs = 10\n    lr = 1e-3\n    batch_size = 128 if USE_TPU else 4\n    embgen_batch_size = batch_size*2 if USE_TPU else batch_size*4\n    num_workers = 2\n    val_samples = 1\n    embedding_size = 256 #128\n    backbone_name = 'efficientnet_b1' #'efficientnet_b0' #'fbnetc_100', 'eca_nfnet_l0'\n    n_classes = data_df[\"hotel_id_code\"].nunique()\n    device = (xm.xla_device() if USE_TPU else ('cuda' if torch.cuda.is_available() else 'cpu'))\n    sct_loss_weight = 0.1\n    lam = 1\n\nprint(args.device)","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:36.742712Z","iopub.execute_input":"2022-06-06T06:59:36.743247Z","iopub.status.idle":"2022-06-06T06:59:36.752176Z","shell.execute_reply.started":"2022-06-06T06:59:36.743210Z","shell.execute_reply":"2022-06-06T06:59:36.751316Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","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","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:42.911046Z","iopub.execute_input":"2022-06-06T06:59:42.911812Z","iopub.status.idle":"2022-06-06T06:59:42.917348Z","shell.execute_reply.started":"2022-06-06T06:59:42.911772Z","shell.execute_reply":"2022-06-06T06:59:42.916565Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset and Transformations","metadata":{}},{"cell_type":"code","source":"import albumentations as A\nimport albumentations.pytorch as APT\nimport cv2 \n\n# used for training dataset - augmentations and occlusions\ntrain_transform = A.Compose([\n    A.HorizontalFlip(p=0.75),\n    A.VerticalFlip(p=0.05), #p=0.25 makes no sense, as there are few vertically-flipped images in the dataset\n    A.ShiftScaleRotate(p=0.5, border_mode=cv2.BORDER_CONSTANT),\n    A.OpticalDistortion(p=0.25),\n    A.Perspective(p=0.25),\n    A.CoarseDropout(p=0.5, min_holes=1, max_holes=6, \n                    min_height=IMG_SIZE//16, max_height=IMG_SIZE//4,\n                    min_width=IMG_SIZE//16,  max_width=IMG_SIZE//4), # normal coarse dropout\n    A.CoarseDropout(p=1., max_holes=1, \n                    min_height=IMG_SIZE//4, max_height=IMG_SIZE//2,\n                    min_width=IMG_SIZE//4,  max_width=IMG_SIZE//2, \n                    fill_value=(255,0,0)),# simulating occlusions in test data\n\n    A.RandomBrightnessContrast(p=0.75),\n    A.ToFloat(),\n    APT.transforms.ToTensorV2(),\n])\n\n# used for validation dataset - only occlusions\nval_transform = A.Compose([\n    A.CoarseDropout(p=1., max_holes=1, \n                    min_height=IMG_SIZE//4, max_height=IMG_SIZE//2,\n                    min_width=IMG_SIZE//4,  max_width=IMG_SIZE//2, \n                    fill_value=(255,0,0)),# simulating occlusions\n    A.ToFloat(),\n    APT.transforms.ToTensorV2(),\n])\n\n# no augmentations\nbase_transform = A.Compose([\n    A.ToFloat(),\n    APT.transforms.ToTensorV2(),\n])","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:45.051832Z","iopub.execute_input":"2022-06-06T06:59:45.052384Z","iopub.status.idle":"2022-06-06T06:59:46.229376Z","shell.execute_reply.started":"2022-06-06T06:59:45.052345Z","shell.execute_reply":"2022-06-06T06:59:46.228668Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","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\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_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        return {\n            \"image\" : image,\n            \"target\" : record['hotel_id_code'],\n        }","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:46.231085Z","iopub.execute_input":"2022-06-06T06:59:46.231408Z","iopub.status.idle":"2022-06-06T06:59:46.238649Z","shell.execute_reply.started":"2022-06-06T06:59:46.231364Z","shell.execute_reply":"2022-06-06T06:59:46.237671Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"# source: https://github.com/ronghuaiyang/arcface-pytorch/blob/master/models/metrics.py\nclass ArcMarginProduct(nn.Module):\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=args.device)\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)\n        output *= self.s\n\n        return output\n\nclass EmbeddingModel(nn.Module):\n    def __init__(self, out_features=100, embed_size=256, backbone_name=\"efficientnet_b1\"):\n        super(EmbeddingModel, 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        self.embedding = nn.Linear(in_features, embed_size)\n        self.classifier = nn.Linear(embed_size, out_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        self.backbone.classifier = nn.Identity()\n\n        self.arc_face = ArcMarginProduct(self.embed_size, out_features, s=30.0, m=0.20, easy_margin=False)\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            nn.BatchNorm1d(self.embed_size),\n        )\n\n        print(f\"Model {backbone_name} ArcMarginProduct - Features: {in_features}, Embeds: {self.embed_size}\")\n        \n    def embed_and_classify(self, x):\n        x = self.forward(x)\n        return x, self.classifier(x)\n    \n    def forward(self, input, targets=None):\n        x = self.backbone(input)\n        x = x.view(x.size(0), -1)\n        x = self.embedding(x) #self.post(x)\n        \n        if targets is not None:\n            logits = self.arc_face(x, targets)\n            return x, logits\n        \n        return x","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:47.243168Z","iopub.execute_input":"2022-06-06T06:59:47.243436Z","iopub.status.idle":"2022-06-06T06:59:47.262962Z","shell.execute_reply.started":"2022-06-06T06:59:47.243407Z","shell.execute_reply":"2022-06-06T06:59:47.262000Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Helper Functions","metadata":{}},{"cell_type":"code","source":"def save_checkpoint(model, scheduler, optimizer, epoch, name, loss=None, score=None):\n    if USE_TPU:\n        checkpoint = {\"epoch\": epoch,\n              \"model\": {k: v.cpu() for k, v in model.state_dict().items()},\n              \"loss\": loss,\n              \"score\": score,\n              }\n    else:\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    os.makedirs(f\"{OUTPUT_FOLDER}checkpoint-{name}\", exist_ok=True)\n    torch.save(checkpoint, f\"{OUTPUT_FOLDER}checkpoint-{name}/epoch{epoch}_score{score}.pt\")\n\ndef load_checkpoint(model, scheduler, optimizer, name):\n    checkpoint = torch.load(name) #torch.load(f\"{OUTPUT_FOLDER}checkpoint-{name}.pt\")\n\n    model.load_state_dict(checkpoint[\"model\"])\n    scheduler.load_state_dict(checkpoint[\"scheduler\"])\n    return model, scheduler, optimizer, checkpoint[\"epoch\"]","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:50.603118Z","iopub.execute_input":"2022-06-06T06:59:50.603456Z","iopub.status.idle":"2022-06-06T06:59:50.614494Z","shell.execute_reply.started":"2022-06-06T06:59:50.603419Z","shell.execute_reply":"2022-06-06T06:59:50.613759Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# method to iterate loader and generate embeddings of images\n# returns embeddings and image class\ndef generate_embeddings(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    targets_all = np.array(targets_all).astype(np.float32)\n    outputs_all = np.array(outputs_all).astype(np.float32)\n            \n    return outputs_all, targets_all","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:51.682091Z","iopub.execute_input":"2022-06-06T06:59:51.682645Z","iopub.status.idle":"2022-06-06T06:59:51.689838Z","shell.execute_reply.started":"2022-06-06T06:59:51.682590Z","shell.execute_reply":"2022-06-06T06:59:51.688949Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training and Validation Functions","metadata":{}},{"cell_type":"code","source":"def train_epoch(args, model, loader, criterion, optimizer, scheduler, epoch, pos_neg_getter):\n    sct_loss_weight = args.sct_loss_weight #(epoch-1)/args.epochs\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(images, targets)\n        loss = criterion(outputs, targets)\n        pos_embeds, neg_embeds = pos_neg_getter(embeds, targets)\n        loss2 = sct_loss(embeds, pos_embeds, neg_embeds)\n        loss = loss+sct_loss_weight*loss2 #(1-sct_loss_weight)*loss+sct_loss_weight*loss2\n        loss.backward()\n        \n        if USE_TPU:\n            xm.optimizer_step(optimizer, barrier=True) #for TPU\n        else:\n            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\"Training epoch {epoch}/{args.epochs} - loss: {loss:0.4f}, accuracy: {score:0.4f}\"\n        t.set_description(desc)\n        \n    return np.mean(losses), score","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:52.817779Z","iopub.execute_input":"2022-06-06T06:59:52.818326Z","iopub.status.idle":"2022-06-06T06:59:52.828516Z","shell.execute_reply.started":"2022-06-06T06:59:52.818290Z","shell.execute_reply":"2022-06-06T06:59:52.827814Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_classification(loader, model):\n    targets_all = []\n    outputs_all = []\n    \n    model.eval()\n    t = tqdm(loader, desc=\"Classification\")\n    \n    for i, sample in enumerate(t):\n        images = sample['image'].to(args.device)\n        targets = sample['target'].to(args.device)\n        \n        _, outputs = model.embed_and_classify(images)\n        \n        targets_all.extend(targets.cpu().numpy())\n        outputs_all.extend(torch.sigmoid(outputs).detach().cpu().numpy())\n        \n    # repeat targets to N_MATCHES for easy calculation of MAP@5\n    y = np.repeat([targets_all], repeats=N_MATCHES, axis=0).T\n    # sort predictions and get top 5\n    preds = np.argsort(-np.array(outputs_all), axis=1)[:, :N_MATCHES]\n    # check if any of top 5 predictions are correct and calculate mean accuracy\n    acc_top_5 = (preds == y).any(axis=1).mean()\n    # calculate prediction accuracy\n    acc_top_1 = np.mean(targets_all == np.argmax(outputs_all, axis=1))\n\n    print(f\"Classification accuracy: {acc_top_1:0.4f}, MAP@5: {acc_top_5:0.4f}\")","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:53.896477Z","iopub.execute_input":"2022-06-06T06:59:53.897388Z","iopub.status.idle":"2022-06-06T06:59:53.905727Z","shell.execute_reply.started":"2022-06-06T06:59:53.897350Z","shell.execute_reply":"2022-06-06T06:59:53.904793Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# find 5 most similar images from different hotels and return their hotel_id_code\ndef find_matches(query, base_embeds, base_targets, k=N_MATCHES):\n    distance_df = pd.DataFrame(index=np.arange(len(base_targets)), data={\"hotel_id_code\": base_targets})\n    # calculate cosine distance of query embeds to all base embeds\n    distance_df[\"distance\"] = cosine_similarity([query], base_embeds)[0]\n    # sort by distance and hotel_id\n    distance_df = distance_df.sort_values(by=[\"distance\", \"hotel_id_code\"], ascending=False).reset_index(drop=True)\n    # return first 5 different hotel_id_codes\n    return distance_df[\"hotel_id_code\"].unique()[:N_MATCHES]\n    \ndef test_similarity(args, data_df, model):#base_loader, test_loader, model):\n    base_embeds = np.stack(data_df[data_df.Set==\"train\"].embeddings.values)\n    base_targets = data_df[data_df.Set==\"train\"].hotel_id_code.values\n    test_embeds = np.stack(data_df[data_df.Set==\"val\"].embeddings.values)\n    test_targets = data_df[data_df.Set==\"val\"].hotel_id_code.values\n    \n    preds = []\n    for query_embeds in tqdm(test_embeds, desc=\"Similarity - match finding\"):\n        tmp = find_matches(query_embeds, base_embeds, base_targets)\n        preds.extend([tmp])\n        \n    preds = np.array(preds)\n    test_targets_N = np.repeat([test_targets], repeats=N_MATCHES, axis=0).T\n    \n    # calculate map@5\n    map_at_5 = (((preds==test_targets_N)*np.array([1.0,1/2,1/3,1/4,1/5])).max(axis=1)).mean()\n    \n    # calculate prediction accuracy\n    acc_top_1 = np.mean(test_targets == preds[:, 0])\n    print(f\"Similarity accuracy: {acc_top_1:0.4f}, MAP@5: {map_at_5:0.4f}\")","metadata":{"execution":{"iopub.status.busy":"2022-06-06T06:59:55.108843Z","iopub.execute_input":"2022-06-06T06:59:55.109380Z","iopub.status.idle":"2022-06-06T06:59:55.121447Z","shell.execute_reply.started":"2022-06-06T06:59:55.109324Z","shell.execute_reply":"2022-06-06T06:59:55.120711Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prepare Data","metadata":{}},{"cell_type":"code","source":"# save hotel_id encoding for later decoding\nhotel_id_code_df = data_df.drop(columns=[\"image_id\"]).drop_duplicates().reset_index(drop=True)\nhotel_id_code_df.to_csv(OUTPUT_FOLDER+'hotel_id_code_mapping.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:00:12.819382Z","iopub.execute_input":"2022-06-06T07:00:12.820057Z","iopub.status.idle":"2022-06-06T07:00:12.840993Z","shell.execute_reply.started":"2022-06-06T07:00:12.820017Z","shell.execute_reply":"2022-06-06T07:00:12.840283Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train and Evaluate","metadata":{}},{"cell_type":"code","source":"model_name = f\"embedding-model-{args.backbone_name}-{IMG_SIZE}x{IMG_SIZE}\"\nprint(model_name)\n\nseed_everything(seed=SEED)","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:00:15.456817Z","iopub.execute_input":"2022-06-06T07:00:15.457271Z","iopub.status.idle":"2022-06-06T07:00:15.466888Z","shell.execute_reply.started":"2022-06-06T07:00:15.457236Z","shell.execute_reply":"2022-06-06T07:00:15.465952Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# split data into train and validation set\nhotel_image_count = data_df.groupby(\"hotel_id\")[\"image_id\"].count()\n# hotels that have more images than samples for validation\nvalid_hotels = hotel_image_count[hotel_image_count > args.val_samples]\n# data that can be split into train and val set\nvalid_data = data_df[data_df[\"hotel_id\"].isin(valid_hotels.index)]\n# if hotel had less than required val_samples it will be only in the train set\n#valid_df = valid_data.groupby(\"hotel_id\").sample(args.val_samples, random_state=SEED)\n#train_df = data_df[~data_df[\"image_id\"].isin(valid_df[\"image_id\"])]\nX_train, X_val, _, _ = train_test_split(valid_data['image_id'], valid_data['hotel_id'], test_size=0.1, stratify=valid_data['hotel_id'])\nvalid_df = data_df[data_df['image_id'].isin(X_val)]\ntrain_df = data_df[data_df['image_id'].isin(X_train)]\n\ntrain_dataset = HotelTrainDataset(train_df, train_transform, data_path=IMAGE_FOLDER)\ntrain_loader  = DataLoader(train_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True, drop_last=True)\nvalid_dataset = HotelTrainDataset(valid_df, val_transform, data_path=IMAGE_FOLDER)\nvalid_loader  = DataLoader(valid_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False)\n# base dataset for image similarity search\nbase_dataset   = HotelTrainDataset(data_df, base_transform, data_path=IMAGE_FOLDER)\nbase_loader    = DataLoader(base_dataset, num_workers=args.num_workers, batch_size=args.embgen_batch_size, shuffle=False)","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:00:31.229245Z","iopub.execute_input":"2022-06-06T07:00:31.229898Z","iopub.status.idle":"2022-06-06T07:00:31.329552Z","shell.execute_reply.started":"2022-06-06T07:00:31.229859Z","shell.execute_reply":"2022-06-06T07:00:31.328815Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = EmbeddingModel(args.n_classes, args.embedding_size ,args.backbone_name)\nmodel = model.to(args.device)","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:00:40.501493Z","iopub.execute_input":"2022-06-06T07:00:40.501978Z","iopub.status.idle":"2022-06-06T07:00:45.771603Z","shell.execute_reply.started":"2022-06-06T07:00:40.501941Z","shell.execute_reply":"2022-06-06T07:00:45.770868Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HardPosHardNegFinder:\n    def __init__(self, compare_targets, args):\n        self.compare_targets = compare_targets.unsqueeze(1)\n        self.inf = torch.tensor(float('inf'), device=args.device)\n    \n    def update_compare_embs(self, compare_embs):\n        self.compare_embs = compare_embs\n        \n    def __call__(self, embeds, targets):\n        sim_combis_mask = targets.unsqueeze(0) == self.compare_targets\n        sim_combis = torch.cosine_similarity(embeds.unsqueeze(0), self.compare_embs.unsqueeze(1), dim=2)\n        same_hotel_most_diff = torch.where(sim_combis_mask, sim_combis, self.inf).argmin(dim=0)        \n        diff_hotel_most_sim = torch.where(~sim_combis_mask, sim_combis, -self.inf).argmax(dim=0)\n\n        return self.compare_embs[same_hotel_most_diff], self.compare_embs[diff_hotel_most_sim]\n\nclass EasyPosHardNegFinder:\n    def __init__(self, compare_targets, args):\n        self.compare_targets = compare_targets.unsqueeze(1)\n        self.ninf = torch.tensor(-float('inf'), device=args.device)\n    \n    def update_compare_embs(self, compare_embs):\n        self.compare_embs = compare_embs\n        \n    def __call__(self, embeds, targets):\n        sim_combis_mask = targets.unsqueeze(0) == self.compare_targets\n        sim_combis = torch.cosine_similarity(embeds.unsqueeze(0), self.compare_embs.unsqueeze(1), dim=2)\n        same_hotel_most_sim = torch.where(sim_combis_mask, sim_combis, self.ninf).argmax(dim=0)\n        diff_hotel_most_sim = torch.where(~sim_combis_mask, sim_combis, self.ninf).argmax(dim=0)\n        return self.compare_embs[same_hotel_most_sim], self.compare_embs[diff_hotel_most_sim]","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:00:45.773830Z","iopub.execute_input":"2022-06-06T07:00:45.774263Z","iopub.status.idle":"2022-06-06T07:00:45.787007Z","shell.execute_reply.started":"2022-06-06T07:00:45.774223Z","shell.execute_reply":"2022-06-06T07:00:45.785831Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def sct_loss(s_an, s_pos, s_neg):\n    S_ap = torch.nn.functional.cosine_similarity(s_an, s_pos)\n    S_an = torch.nn.functional.cosine_similarity(s_an, s_neg)\n    return torch.where(S_ap>S_an, args.lam*S_an, triplet_loss(S_ap, S_an)).mean()\n\ndef triplet_loss(S_ap, S_an):\n    e_S_ap = torch.exp(S_ap)\n    e_S_an = torch.exp(S_an)\n    return -torch.log(torch.div(e_S_ap, e_S_ap+e_S_an))","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:00:49.860795Z","iopub.execute_input":"2022-06-06T07:00:49.861067Z","iopub.status.idle":"2022-06-06T07:00:49.866917Z","shell.execute_reply.started":"2022-06-06T07:00:49.861034Z","shell.execute_reply":"2022-06-06T07:00:49.866181Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sets_df = pd.concat([\n    train_df.set_index('image_id').assign(Set='train').Set,\n    valid_df.set_index('image_id').assign(Set='val').Set\n])\ndata_df = data_df.join(sets_df, on='image_id')","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:00:50.710464Z","iopub.execute_input":"2022-06-06T07:00:50.711176Z","iopub.status.idle":"2022-06-06T07:00:50.747054Z","shell.execute_reply.started":"2022-06-06T07:00:50.711138Z","shell.execute_reply":"2022-06-06T07:00:50.746252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"criterion = nn.CrossEntropyLoss()\noptimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)\n\nscheduler = torch.optim.lr_scheduler.OneCycleLR(\n                optimizer,\n                max_lr=args.lr,\n                epochs=args.epochs, #20\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\nstart_epoch = 1","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:00:52.518555Z","iopub.execute_input":"2022-06-06T07:00:52.519293Z","iopub.status.idle":"2022-06-06T07:00:52.526944Z","shell.execute_reply.started":"2022-06-06T07:00:52.519258Z","shell.execute_reply":"2022-06-06T07:00:52.526149Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs(f\"{OUTPUT_FOLDER}{model_name}_image-embeddings\", exist_ok=True)\n\nbase_embeds, _ = generate_embeddings(base_loader, model, \"Generate embeddings for all images\")\ndata_df[\"embeddings\"] = list(base_embeds)\ndata_df.to_pickle(f\"{OUTPUT_FOLDER}{model_name}_image-embeddings/epoch0.pkl\")\n#data_df.to_pickle(f\"{OUTPUT_FOLDER}{model_name}_image-embeddings/epoch{start_epoch-1}.pkl\")","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:00:53.932844Z","iopub.execute_input":"2022-06-06T07:00:53.933636Z","iopub.status.idle":"2022-06-06T07:10:12.841894Z","shell.execute_reply.started":"2022-06-06T07:00:53.933583Z","shell.execute_reply":"2022-06-06T07:10:12.841045Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_similarity(args, data_df, model)","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:10:12.846899Z","iopub.execute_input":"2022-06-06T07:10:12.849796Z","iopub.status.idle":"2022-06-06T07:18:16.737021Z","shell.execute_reply.started":"2022-06-06T07:10:12.849754Z","shell.execute_reply":"2022-06-06T07:18:16.736286Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_hotel_id_code = torch.from_numpy(data_df[(data_df.Set=='train')].hotel_id_code.values).to(args.device)\ntrain_hotel_id_code.requires_grad = False\n\ntrain_embs = torch.from_numpy(np.stack(data_df[(data_df.Set=='train')].embeddings.values)).to(args.device)\ntrain_embs.requires_grad = False\n\n#hard_pos_hard_neg_getter = HardPosHardNegFinder(train_hotel_id_code, args)\npos_neg_getter = EasyPosHardNegFinder(train_hotel_id_code, args)\npos_neg_getter.update_compare_embs(train_embs)","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:18:16.738552Z","iopub.execute_input":"2022-06-06T07:18:16.739080Z","iopub.status.idle":"2022-06-06T07:18:16.871838Z","shell.execute_reply.started":"2022-06-06T07:18:16.739040Z","shell.execute_reply":"2022-06-06T07:18:16.871039Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for epoch in range(start_epoch, args.epochs+1):\n    train_loss, train_score = train_epoch(args, model, train_loader, criterion, optimizer, scheduler, epoch, pos_neg_getter)\n    save_checkpoint(model, scheduler, optimizer, epoch, model_name, train_loss, train_score)\n    \n    # generate embeddings for all train images and save them for inference\n    base_embeds, _ = generate_embeddings(base_loader, model, \"Generate embeddings for all images\")\n    data_df[\"embeddings\"] = list(base_embeds)\n    data_df.to_pickle(f\"{OUTPUT_FOLDER}{model_name}_image-embeddings/epoch{epoch}.pkl\")\n    test_similarity(args, data_df, model)\n    train_embs = torch.from_numpy(np.stack(data_df[(data_df.Set=='train')].embeddings.values)).to(args.device)\n    train_embs.requires_grad = False\n    pos_neg_getter.update_compare_embs(train_embs)","metadata":{"execution":{"iopub.status.busy":"2022-06-06T07:18:16.873706Z","iopub.execute_input":"2022-06-06T07:18:16.873964Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"base_dataset  = HotelTrainDataset(train_df, base_transform, data_path=IMAGE_FOLDER)\nbase_loader   = DataLoader(base_dataset, num_workers=args.num_workers, batch_size=args.embgen_batch_size, shuffle=False)\ntest_similarity(args, data_df, model)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}