{"cells":[{"metadata":{},"cell_type":"markdown","source":"The hope of this notebook was to get the final layer of the convolutional network as an embedding. By using the arcface loss I hoped that the embeddings would be represent the relevant cluster, hence rather than taking the cluster center, any random example of a class should have similar embeddings to other members of a class. However, this doesn't seem to be the case.\n\nSee version 6 of this notebook to see actual results.\n\n## What's Special about this NB:\n- Used Nvidia Apex to get mixed precision on top of pytorch\n- Used Lamb optimizer + lr scheduler which decays learning rate every 100 iterations\n- Used ArcFace loss. Jump to `Loss function` section to see calculation. It is different from the author's definition.\n- Used efficientnet as base\n- Only trained the weight centers and the final convolutional layer of efficientnet.\n- 30 minutes to run through 1 epoch of ~300k images (since I only took any class that had >100 instances). See dataset and DataLoader section to see what I did to get it to run this fast.\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%capture\n!git clone https://github.com/NVIDIA/apex\n!pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./apex\n!pip install efficientnet_pytorch torchtoolbox","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nimport random\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom tqdm.auto import tqdm\nimport multiprocessing as mp\nfrom functools import partial\nimport pickle\n\nimport cv2\n\nfrom sklearn.model_selection import train_test_split\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import _LRScheduler, StepLR\n\ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n\nfrom apex.optimizers import FusedAdam, FusedLAMB\nfrom apex import amp\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom efficientnet_pytorch import EfficientNet\n\nimport warnings\nwarnings.simplefilter('ignore')\n%matplotlib inline\n\ntorch.backends.cudnn.benchmark = True","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Config"},{"metadata":{"trusted":true},"cell_type":"code","source":"BASE = \"efficientnet-b0\"\nSIZE = (128, 128)\nEPOCHS = 2\nGRAD_ACCUMULATE = 1\nBS = 512\np = 0.5\nLR_RANGE = [1e-7, 2e-4]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data"},{"metadata":{},"cell_type":"markdown","source":"### Get Training/ Validation dataframes"},{"metadata":{"trusted":true},"cell_type":"code","source":"path = \"/kaggle/input/landmark-recognition-2020/\"\ndf = pd.read_csv(path + \"train.csv\")\ndf[\"path\"] = df[\"id\"].map(lambda x: \"/\".join([path+\"train\"] + list(x[:3])+[x + \".jpg\"]))\ndf.sort_values(\"landmark_id\", inplace=True)\ndf.reset_index(drop=True, inplace=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"k = 100\ncounts = df[\"landmark_id\"].value_counts()\ntopk = counts[counts >= k].index\ndf = df[df[\"landmark_id\"].isin(topk)]\ndf.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"id2y = {id_:i for i, id_ in enumerate(df[\"landmark_id\"].unique())}\ndf[\"target\"] = df[\"landmark_id\"].map(lambda x: id2y[x])\n\nweights = 1 / df[\"target\"].value_counts()\ndf[\"weights\"] = weights.loc[df[\"target\"]].values\nprint(f\"There are {len(id2y)} classes\")\n\nprint(df.shape)\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df, val_df = train_test_split(df, test_size=0.15, random_state=42, stratify=df[\"target\"])\ntrain_df.reset_index(drop=True, inplace=True)\ntrain_df = train_df.sample(frac=1)\nval_df.reset_index(drop=True, inplace=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_files = []\nfor root, dirs, files in tqdm(os.walk(path+\"test/\")):\n    if files:\n        files = [root+\"/\"+file for file in files]\n        test_files.extend(files)\n        \ntest_df = pd.DataFrame({\"path\": test_files})","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### PyTorch Datasets + DataLoaders"},{"metadata":{"trusted":true},"cell_type":"code","source":"tfms = A.Compose([\n    A.Normalize(\n        mean=[0.485, 0.456, 0.406],\n        std=[0.229, 0.224, 0.225],\n    ),\n    ToTensorV2()\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Images(Dataset):\n    def __init__(self, df: pd.DataFrame, train: bool = True):\n        \"\"\"\n        Parameters:\n            df (pd.DataFrame): DataFrame with data description\n            train (bool): flag of whether a training dataset is being initialized or testing one\n            transforms: image transformation method to be applied\n        \"\"\"\n        if train:\n            df[\"weights\"] = df[\"weights\"].astype(np.float32)\n        self.df = df.reset_index(drop=True)\n        self.train = train\n        \n    def __getitem__(self, index):\n        im_path = self.df.loc[index, 'path']\n        x = cv2.cvtColor(cv2.resize(cv2.imread(im_path), SIZE), cv2.COLOR_BGR2RGB)\n        x = tfms(image=x)['image']\n            \n        if self.train:\n            weights = self.df.loc[index, 'weights']\n            y = self.df.loc[index, 'target']\n            return x, y, weights\n        else:\n            return x\n    \n    def __len__(self):\n        return len(self.df)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_images = Images(train_df)\nval_images = Images(val_df)\ntest_images = Images(test_df, train=False)\n\ntrain_dl = DataLoader(train_images, BS, num_workers=mp.cpu_count(), pin_memory=True, shuffle=True, drop_last=True)\nval_dl = DataLoader(val_images, BS, num_workers=mp.cpu_count(), pin_memory=True)\ntest_dl = DataLoader(test_images, BS*4, num_workers=mp.cpu_count(), pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"class ArcFaceLoss(nn.Module):\n    def __init__(self, s=64, m=0.5):\n        super().__init__()\n        self.s, self.m = s, m\n        self.cross_entropy = partial(F.cross_entropy, reduction='none')\n        \n    def forward(self, costheta, y):\n        costheta_y = costheta[torch.arange(len(y)), y]\n        costheta_y = torch.cos(torch.acos(costheta_y) +self.m)\n        costheta[torch.arange(len(y)), y] = costheta_y.type(costheta.dtype)\n        \n        return self.cross_entropy(self.s*costheta, y)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Model(nn.Module):\n    def __init__(self, classes, base=BASE, unfreeze=None):\n        super().__init__()\n\n        # EfficientNet\n        self.base = EfficientNet.from_pretrained(base)\n                \n        # Replace last layer\n        self.centers = nn.Parameter(torch.randn(self.base._fc.in_features, classes))\n        self.unfreeze = unfreeze\n    \n    def get_embedding(self, x):\n        pool = F.adaptive_avg_pool2d(self.base.extract_features(x), 1)\n        pool = pool.view(x.shape[0], -1)\n        \n        lens = torch.sqrt((pool**2).sum(dim=-1, keepdim=True))\n        return pool / lens\n    \n    def forward(self, x):\n        embeds = self.get_embedding(x)\n        \n        lens = torch.sqrt((self.centers**2).sum(dim=0, keepdim=True))\n        centers = self.centers / lens\n        \n        return embeds.matmul(centers)\n    \n    def freeze(self):\n        for n,p in self.named_parameters():\n            if not any([layer in n for layer in self.unfreeze]):\n                p.requires_grad = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"unfreeze = [\"conv_head\", \"base._bn1\", \"centers\"]\nmodel = Model(len(id2y), unfreeze=unfreeze)\nmodel = model.to(device)\nmodel.freeze()\n\narcface_loss = ArcFaceLoss()\n\noptimizer = FusedLAMB(model.parameters(), lr=1e-2)\nmodel, optimizer = amp.initialize(model, optimizer, opt_level=\"O1\")\nscheduler = StepLR(optimizer, step_size=100, gamma=0.7)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def metrics(y_pred, y, weights, k=5):\n    \"\"\"\n    Weighted accuracy and top-k accuracy\n    parameters:\n    - y_pred: predicted logits or probabilities \n    - y: Actual class\n    - weights: importance of each instance **must sum to one**\n    - k: number of categories to look for\n    \"\"\"\n    topk = y_pred.topk(k=k, dim=-1)[1] == y[:, None]\n    topk_acc = (weights * topk.any(dim=-1).float()).sum()\n    acc = (weights * topk[:,0].float()).sum()\n    return acc, topk_acc","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"val_p = 0.5\nval_k = int(val_p * len(train_dl))\nval_accs = []\nval_topk_accs = []\n\naccs = []\ntopk_accs = []\nlosses = []\n\nfor p in model.parameters(): p.grad = None\nmodel.train()\nfor _ in range(EPOCHS):\n    for i, (x, y, weights) in tqdm(enumerate(train_dl), total=len(train_dl)):\n        x, y, weights = x.to(device), y.to(device), weights.to(device)\n\n        y_pred = model(x)\n        loss_all = arcface_loss(y_pred, y)\n        loss = (loss_all * weights).sum()\n        loss = loss / GRAD_ACCUMULATE\n\n        acc, topk_acc = metrics(y_pred, y, weights / weights.sum())\n        print(f\"\\rLoss: {loss:.4f}, Accuracy: {acc:.4f}, top-5 Accuracy: {topk_acc:.4f}\", end=\"\")\n\n        with amp.scale_loss(loss, optimizer) as scaled_loss:\n            scaled_loss.backward()\n\n        if (i+1) % GRAD_ACCUMULATE == 0:\n            optimizer.step()\n            for p in model.parameters(): p.grad = None\n            scheduler.step()\n\n        losses.append(loss.detach().cpu().numpy())\n        accs.append(acc)\n        topk_accs.append(topk_acc)\n\n        if (i+1) % val_k == 0:\n            print(\"\\n\")\n            sum_weights = 0.0\n            loss = 0.0\n            for x, y, weights in tqdm(val_dl):\n                x, y, weights = x.to(device), y.to(device), weights.to(device)\n                sum_weights += weights.sum()\n                with torch.no_grad():\n                    model.eval()\n                    y_pred = model(x)\n                    loss_all = arcface_loss(y_pred, y)\n                    loss += (loss_all * weights).sum()\n                    acc, topk_acc = metrics(y_pred, y, weights/weights.sum())\n                    val_accs.append(acc)\n                    val_topk_accs.append(topk_acc)\n            print(f\"\\nValidation Loss: {loss / sum_weights:.4f}, Accuracy: {acc:.4f}, top-5 Accuracy: {topk_acc:.4f}\\n\")\n            model.train()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"topk_accs[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"accs = torch.stack(accs).cpu().numpy()\ntopk_accs = torch.stack(topk_accs).cpu().numpy()\nval_accs = torch.stack(val_accs).cpu().numpy()\nval_topk_accs = torch.stack(val_topk_accs).cpu().numpy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(losses)\nplt.title(\"Loss\")\nplt.show()\nplt.plot(accs)\nplt.title(\"Accuracy\")\nplt.show()\nplt.plot(topk_accs)\nplt.title(\"Top-k\")\nplt.show()\n\nplt.plot(val_accs)\nplt.title(\"Validation Accuracy\")\nplt.show()\nplt.plot(val_topk_accs)\nplt.title(\"Validation Top-k\")\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Save"},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model.state_dict(), \"model.ckpt\")\nwith open(\"id2y.pickle\", \"wb\") as f:\n    pickle.dump(id2y, f)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Evaluate"},{"metadata":{"trusted":true},"cell_type":"code","source":"path = \"/kaggle/input/landmark-recognition-2020/\"\ndf = pd.read_csv(path + \"train.csv\")\ndf[\"path\"] = df[\"id\"].map(lambda x: \"/\".join([path+\"train\"] + list(x[:3])+[x + \".jpg\"]))\ndf.sort_values(\"landmark_id\", inplace=True)\ndf.reset_index(drop=True, inplace=True)\n\ngrp_df = df.groupby(\"landmark_id\").head(1)\nint2id = {i:id_ for i, id_ in enumerate(grp_df[\"landmark_id\"].values)}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"grp_images = Images(grp_df, train=False)\ntest_images = Images(test_df, train=False)\nBS = 1024\n\ngrp_dl = DataLoader(grp_images, BS, num_workers=mp.cpu_count(), pin_memory=True)\ntest_dl = DataLoader(test_images, BS, num_workers=mp.cpu_count(), pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_embeds = []\nwith torch.no_grad():\n    model.eval()\n    for x in tqdm(grp_dl):\n        x = x.to(device)\n        train_embeds.extend(model.get_embedding(x))\n\ntrain_embeds = torch.stack(train_embeds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ps = []\ncategories = []\nwith torch.no_grad():\n    model.eval()\n    for x in tqdm(test_dl):\n        x = x.to(device)\n        test = model.get_embedding(x)\n        p, category = train_embeds.matmul(test.T).max(dim=0)\n        category = [int2id[i] for i in category.cpu().numpy()]\n        ps.extend((p + 1) * 0.5)\n        categories.extend(category)\n\nps = torch.stack(ps).cpu().numpy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df[\"id\"] = test_df[\"path\"].map(lambda x: os.path.basename(x)[:-4])\ntest_df[\"landmark\"] = categories\ntest_df[\"p\"] = ps\ntest_df[\"landmarks\"] = test_df[\"landmark\"].astype(str) + test_df[\"p\"].map(lambda x: f\" {x:.4f}\")\ncols = [\"id\", \"landmarks\"]\ntest_df[cols].to_csv(\"submission.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = test_df.sample(frac=1)\nfor i in range(5):\n    plt.subplot(121)\n    im_path = test_df[\"path\"].values[i]\n    plt.imshow(cv2.cvtColor(cv2.resize(cv2.imread(im_path), SIZE), cv2.COLOR_BGR2RGB))\n    plt.subplot(122)\n    im_path = grp_df.loc[grp_df[\"landmark_id\"] == test_df[\"landmark\"].values[i], \"path\"].values[0]\n    plt.imshow(cv2.cvtColor(cv2.resize(cv2.imread(im_path), SIZE), cv2.COLOR_BGR2RGB))\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}