{"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":"# ArcFaces and EfficientNet with Pytorch Lightning ⚡\nHi Kagglers ! 👋<br>\nIn this notebook I want to make sort of a baseline with pytorch lightning as I believe this library is very helpful and can help you iterate quickly.<br>\nOn a high level, this notebook translates into pytorch and merges the following solutions:\n1. The template that is used by best public notebooks\n2. The 3rd place solution of the [Google Landmark Recognition 2020](https://www.kaggle.com/c/landmark-recognition-2020)\n3. The current best public pytorch notebook [Pytorch inference notebok (Arcface + GeM Pooling)](https://www.kaggle.com/vladvdv/pytorch-inference-notebok-arcface-gem-pooling)\n\n## What is implemented in this notebook ?\n# What methods are used in this notebook 💡 ?\n\n| Method | Learn more | Description |\n|:---|:---|:---|\n| EfficientNet | [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) | Architecture designed for efficiency. As we have limited GPU quota, I want to be able to run the best model on a limited setup |\n| ArcFaces sub-centers | [Sub-center ArcFace: Boosting Face Recognition by Large-scale Noisy Web Faces](https://paperswithcode.com/paper/sub-center-arcface-boosting-face-recognition) | Each class have multiple embeddings and you are assigning to each class the distance of its clothest sub-center |\n| ArcFaces Loss | [ArcFace: Additive Angular Margin Loss for Deep Face Recognition](https://arxiv.org/abs/1801.07698) | Making it harder for the model to predict the right class during training as a fegularization method |\n| Dynamic margins | [Google Landmark Recognition 2020 Competition Third Place Solution](https://arxiv.org/abs/2010.05350) | Create a specific margin for each class based on its frequency in the dataset |\n| GeM Pooling | [Fine-tuning CNN Image Retrieval with No Human Annotation](https://arxiv.org/abs/1711.02512) | Pooling method sort of intermediate between max-pooling and average-pooling. How much is it similar to one or the other is learned |\n| Gradient accumulation | [Accumulate gradients](https://pytorch-lightning.readthedocs.io/en/stable/advanced/training_tricks.html#accumulate-gradients) | Accumulate gradient over multiple batches to simulate any batch size regardless of your GPU memory size |\n| Mixed precision training | [Precision](https://pytorch-lightning.readthedocs.io/en/latest/advanced/mixed_precision.html) | Train models using 16 bits data types to be able to fit models you couldn't on your hardware without it |","metadata":{}},{"cell_type":"markdown","source":"# 📚 Librairies ","metadata":{}},{"cell_type":"code","source":"!export CUDA_LAUNCH_BLOCKING=1\nimport numpy as np\nimport pandas as pd\nimport os\nimport torch\nimport random\nimport albumentations\nimport cv2\nimport math\nimport copy\nimport torchmetrics\nimport torchvision\nimport pytorch_lightning as pl\nimport wandb\nimport json\nfrom torch import nn, Tensor\nimport torch.nn.functional as F\nfrom sklearn.model_selection import train_test_split\nfrom torch.utils.data import Dataset\nimport matplotlib.pyplot as plt\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom torchvision import transforms, models\nfrom pytorch_lightning import Callback, LightningModule, Trainer\nfrom pytorch_lightning.loggers import WandbLogger\nfrom sklearn.neighbors import NearestNeighbors\n!pip install timm\nimport timm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-03-07T18:25:13.018152Z","iopub.execute_input":"2022-03-07T18:25:13.018589Z","iopub.status.idle":"2022-03-07T18:25:29.307756Z","shell.execute_reply.started":"2022-03-07T18:25:13.018446Z","shell.execute_reply":"2022-03-07T18:25:29.306921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configuration ⚙️\nFeel free to fork and change the configuration here:","metadata":{}},{"cell_type":"code","source":"class CFG:\n    SEED = 69\n    ### Dataset\n    ## Effective batch size will be BATCH_SIZE*ACCUMULATE_GRAD_BATCHES\n    #BATCH_SIZE = 16\n    BATCH_SIZE = 48\n    #ACCUMULATE_GRAD_BATCHES = 8\n    ACCUMULATE_GRAD_BATCHES = 1\n    IMAGE_SIZE = 512\n    #IMAGE_SIZE = 224\n    NUM_WORKERS = 2\n    ### Model\n    MODEL_NAME = \"efficientnet_b0\"\n    EMBEDDING_SIZE = 512\n    ### Training\n    ## Arcfaces\n    CENTERS_PER_CLASS = 3\n    S = 30\n    MARGIN_MIN = 0.2\n    MARGIN_MAX = 0.4\n    EPOCHS = 20\n    MIXED_PRECISION = True\n    MODEL_PATH=\"model.ckpt\"\n    # Inference\n    KNN = 100\n    Q_NEW = 0.112 # Proportion of new individuals expected in the dataset","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:29.311227Z","iopub.execute_input":"2022-03-07T18:25:29.311466Z","iopub.status.idle":"2022-03-07T18:25:29.317180Z","shell.execute_reply.started":"2022-03-07T18:25:29.311435Z","shell.execute_reply":"2022-03-07T18:25:29.316515Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Logging 📄\nLogging using Weights and Biases 🪄🐝","metadata":{}},{"cell_type":"code","source":"from kaggle_secrets import UserSecretsClient\n\nuser_secrets = UserSecretsClient()\n\n# I have saved my API token with \"wandb_api\" as Label. \n# If you use some other Label make sure to change the same below. \nwandb_api = user_secrets.get_secret(\"wandb_api\") \n\nwandb.login(key=wandb_api)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:29.318507Z","iopub.execute_input":"2022-03-07T18:25:29.318969Z","iopub.status.idle":"2022-03-07T18:25:31.261084Z","shell.execute_reply.started":"2022-03-07T18:25:29.318931Z","shell.execute_reply":"2022-03-07T18:25:31.260270Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Make a deterministic pipeline\ndef fix_seed(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    #gpu randomseed fixed\n    torch.backends.cudnn.deterministic = True\n\nfix_seed(CFG.SEED)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:31.263291Z","iopub.execute_input":"2022-03-07T18:25:31.263743Z","iopub.status.idle":"2022-03-07T18:25:31.272187Z","shell.execute_reply.started":"2022-03-07T18:25:31.263698Z","shell.execute_reply":"2022-03-07T18:25:31.271348Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset 🖼️\nImplementing the dataset as a Pytorch Dataset as required by Pytorch Lightning. It applies some augmentations:","metadata":{}},{"cell_type":"code","source":"BASE_PATH = \"../input/happywhale-enhanced-dataset-large\"\n\nclass WandDID(Dataset):\n    def __init__(self, data, augment=True, base_path=BASE_PATH):\n        self.base_path = os.path.join(base_path, \"train_images\")\n        self.data = data\n        # Augmentations\n        if augment:\n            transformations = albumentations.Compose([\n                albumentations.augmentations.transforms.HueSaturationValue(p=0.5),\n                albumentations.Cutout(num_holes=1, max_h_size=(CFG.IMAGE_SIZE//9), max_w_size=(CFG.IMAGE_SIZE//9), p=0.5),\n                albumentations.RandomBrightnessContrast(brightness_limit=0.10, contrast_limit=(-0.2, 0.2),p=0.5),\n                albumentations.Normalize(),\n                ToTensorV2(p=1.0)\n            ])\n        else:\n            transformations = albumentations.Compose([\n                albumentations.Normalize(),\n                ToTensorV2(p=1.0)\n            ])\n\n        def make_transform(transform=False):\n            def f(image):\n                if transform:\n                    image_np = np.array(image)\n                    augmented = transform(image=image_np)\n                return augmented\n            return f\n\n        self.transforms = transforms.Compose([\n            transforms.Lambda(make_transform(transformations)),\n        ])\n        \n    def __getitem__(self, idx):\n        image = self.preprocess(self.data[\"image\"].iloc[idx])\n        label = self.data[\"individual_id_integer\"].iloc[idx]\n        return image, torch.tensor(label, dtype=torch.long)\n    \n    def preprocess(self, image):\n        image = os.path.join(self.base_path, image)\n        image = cv2.imread(image)[:, :, ::-1]\n        if image.shape[0] < CFG.IMAGE_SIZE or image.shape[1] < CFG.IMAGE_SIZE:\n            image = cv2.resize(image, (CFG.IMAGE_SIZE, CFG.IMAGE_SIZE), cv2.INTER_CUBIC)\n        if self.transforms is not None:\n            image = self.transforms(image)[\"image\"]\n        return image\n    \n    def plot_sample(self, idx):\n        image = self.data[\"image\"].iloc[idx]\n        image = os.path.join(self.base_path, image)\n        image = cv2.imread(image)[:, :, ::-1]\n        plt.title(\"{} ({})\".format(\n            self.data[\"individual_id\"].iloc[idx],\n            self.data[\"species\"].iloc[idx]\n        ))\n        plt.imshow(image)\n        plt.show()\n    \n    def __len__(self):\n        return len(self.data)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:31.273997Z","iopub.execute_input":"2022-03-07T18:25:31.274836Z","iopub.status.idle":"2022-03-07T18:25:31.292588Z","shell.execute_reply.started":"2022-03-07T18:25:31.274794Z","shell.execute_reply":"2022-03-07T18:25:31.291864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data = pd.read_csv(os.path.join(BASE_PATH, \"train.csv\"))\nN_CLASSES = len(data[\"individual_id\"].unique())\n# Computing an integer mapping for individuals ids\nindividual_mapping = {k:i for i, k in enumerate(data[\"individual_id\"].unique())}\n# Compute margins for ArcFaces with dynamic margins\ntmp = np.sqrt(1 / np.sqrt(data['individual_id'].value_counts().loc[list(individual_mapping)].values))\nMARGINS = (tmp - tmp.min()) / (tmp.max() - tmp.min()) * (CFG.MARGIN_MAX - CFG.MARGIN_MIN) + CFG.MARGIN_MIN\n# Save individual mapping\nwith open(\"individual_mapping.json\", \"w\") as f:\n    json.dump(individual_mapping, f)\n\ndata[\"individual_id_integer\"] = data[\"individual_id\"].map(individual_mapping)\ntrain_dataset = WandDID(data)\n# Dataloader\ntrain_loader = torch.utils.data.DataLoader(\n    train_dataset,\n    batch_size=CFG.BATCH_SIZE,\n    num_workers=CFG.NUM_WORKERS,\n    pin_memory=True,\n    shuffle=True\n)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:31.294133Z","iopub.execute_input":"2022-03-07T18:25:31.295086Z","iopub.status.idle":"2022-03-07T18:25:31.503991Z","shell.execute_reply.started":"2022-03-07T18:25:31.295044Z","shell.execute_reply":"2022-03-07T18:25:31.501486Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset.plot_sample(random.randint(0, len(train_dataset)-1))","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:31.509551Z","iopub.execute_input":"2022-03-07T18:25:31.512154Z","iopub.status.idle":"2022-03-07T18:25:31.948621Z","shell.execute_reply.started":"2022-03-07T18:25:31.512103Z","shell.execute_reply":"2022-03-07T18:25:31.947862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model 🤖\nImplements the model as a Pytorch lightning module. We use EfficientNetB4 as feature extractor, the upper layers of the 3rd google landmark solution, and GeM Pooling","metadata":{}},{"cell_type":"code","source":"\"\"\"\n    https://github.com/haqishen/Google-Landmark-Recognition-2020-3rd-Place-Solution\n\"\"\"\nclass Swish(torch.autograd.Function):\n\n    @staticmethod\n    def forward(ctx, i):\n        result = i * torch.sigmoid(i)\n        ctx.save_for_backward(i)\n        return result\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        i = ctx.saved_variables[0]\n        sigmoid_i = torch.sigmoid(i)\n        return grad_output * (sigmoid_i * (1 + i * (1 - sigmoid_i)))\n\n\nclass Swish_module(nn.Module):\n    def forward(self, x):\n        return Swish.apply(x)\n    \nclass DenseCrossEntropy(nn.Module):\n    def forward(self, x, target):\n        x = x.float()\n        target = target.float()\n        logprobs = torch.nn.functional.log_softmax(x, dim=-1)\n\n        loss = -logprobs * target\n        loss = loss.sum(-1)\n        return loss.mean()\n\nclass ArcMarginProduct_subcenter(nn.Module):\n    def __init__(self, in_features=CFG.EMBEDDING_SIZE, out_features=N_CLASSES, k=CFG.CENTERS_PER_CLASS):\n        super().__init__()\n        self.weight = nn.Parameter(torch.FloatTensor(out_features*k, in_features))\n        self.reset_parameters()\n        self.k = k\n        self.out_features = out_features\n        \n    def reset_parameters(self):\n        stdv = 1. / math.sqrt(self.weight.size(1))\n        self.weight.data.uniform_(-stdv, stdv)\n        \n    def forward(self, features):\n        cosine_all = F.linear(F.normalize(features), F.normalize(self.weight))\n        cosine_all = cosine_all.view(-1, self.out_features, self.k)\n        cosine, _ = torch.max(cosine_all, dim=2)\n        return cosine   \n    \nclass ArcFaceLossAdaptiveMargin(nn.modules.Module):\n    def __init__(self, margins, out_dim=N_CLASSES, s=CFG.S):\n        super().__init__()\n        self.crit = DenseCrossEntropy()\n        self.s = s\n        self.register_buffer('margins', torch.tensor(margins))\n        self.out_dim = out_dim\n            \n    def forward(self, logits, labels):\n        #ms = []\n        #ms = self.margins[labels.cpu().numpy()]\n        ms = self.margins[labels]\n        cos_m = torch.cos(ms) #torch.from_numpy(np.cos(ms)).float().cuda()\n        sin_m = torch.sin(ms) #torch.from_numpy(np.sin(ms)).float().cuda()\n        th = torch.cos(math.pi - ms)#torch.from_numpy(np.cos(math.pi - ms)).float().cuda()\n        mm = torch.sin(math.pi - ms) * ms#torch.from_numpy(np.sin(math.pi - ms) * ms).float().cuda()\n        labels = F.one_hot(labels, self.out_dim)\n        labels = labels.half() if CFG.MIXED_PRECISION else labels.float()\n        cosine = logits\n        sine = torch.sqrt(1.0 - cosine * cosine)\n        phi = cosine * cos_m.view(-1,1) - sine * sin_m.view(-1,1)\n        phi = torch.where(cosine > th.view(-1,1), phi, cosine - mm.view(-1,1))\n        output = (labels * phi) + ((1.0 - labels) * cosine)\n        output *= self.s\n        loss = self.crit(output, labels)\n        return loss","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:31.950263Z","iopub.execute_input":"2022-03-07T18:25:31.950790Z","iopub.status.idle":"2022-03-07T18:25:31.987302Z","shell.execute_reply.started":"2022-03-07T18:25:31.950745Z","shell.execute_reply":"2022-03-07T18:25:31.986549Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"\n    https://www.kaggle.com/vladvdv/pytorch-inference-notebok-arcface-gem-pooling\n\"\"\"\n\nclass GeM(nn.Module):\n    def __init__(self, p=3, eps=1e-6):\n        super(GeM, self).__init__()\n        self.p = nn.Parameter(torch.ones(1)*p)\n        self.eps = eps\n\n    def forward(self, x):\n        return self.gem(x, p=self.p, eps=self.eps)\n        \n    def gem(self, x, p=3, eps=1e-6):\n        return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1./p)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:31.993641Z","iopub.execute_input":"2022-03-07T18:25:31.996282Z","iopub.status.idle":"2022-03-07T18:25:32.006592Z","shell.execute_reply.started":"2022-03-07T18:25:31.996240Z","shell.execute_reply":"2022-03-07T18:25:32.005978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":" \"\"\"\n     Scheduler coming from: https://www.kaggle.com/andrej0marinchenko/happywhale-0-679\n \"\"\"\n    \nclass PeakScheduler(torch.optim.lr_scheduler._LRScheduler):\n        def __init__(\n                self, optimizer,\n                epoch_size=-1,\n                lr_start   = 0.000001,\n                lr_max     = 0.000005 * CFG.BATCH_SIZE * CFG.ACCUMULATE_GRAD_BATCHES,\n                lr_min     = 0.000001,\n                lr_ramp_ep = 4,\n                lr_sus_ep  = 0,\n                lr_decay   = 0.8,\n                verbose = True\n            ):\n            self.epoch_size = epoch_size\n            self.optimizer= optimizer\n            self.lr_start = lr_start\n            self.lr_max = lr_max\n            self.lr_min = lr_min\n            self.lr_ramp_ep = lr_ramp_ep\n            self.lr_sus_ep = lr_sus_ep\n            self.lr_decay = lr_decay\n            self.is_plotting = True\n            epochs = list(range(CFG.EPOCHS))\n            learning_rates = []\n            for i in epochs:\n                self.epoch = i\n                learning_rates.append(self.get_lr())\n            self.is_plotting = False\n            self.epoch = 0\n            plt.scatter(epochs,learning_rates)\n            plt.show()\n            super(PeakScheduler, self).__init__(optimizer, verbose=verbose)\n\n        def get_lr(self):\n            if not self.is_plotting:\n                if self.epoch_size == -1:\n                    self.epoch = self._step_count - 1\n                else:\n                    self.epoch = (self._step_count - 1) / self.epoch_size\n                    \n            if self.epoch < self.lr_ramp_ep:\n                lr = (self.lr_max - self.lr_start) / self.lr_ramp_ep * self.epoch + self.lr_start\n\n            elif self.epoch < self.lr_ramp_ep + self.lr_sus_ep:\n                lr = self.lr_max\n            else:\n                lr = (self.lr_max - self.lr_min) * self.lr_decay**(self.epoch - self.lr_ramp_ep - self.lr_sus_ep) + self.lr_min\n            return [lr for _ in self.optimizer.param_groups]","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:32.013210Z","iopub.execute_input":"2022-03-07T18:25:32.013650Z","iopub.status.idle":"2022-03-07T18:25:32.030429Z","shell.execute_reply.started":"2022-03-07T18:25:32.013613Z","shell.execute_reply":"2022-03-07T18:25:32.029564Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class WandDIDNet(LightningModule):\n    def __init__(self, continuous_scheduler=True, s=30, m=0.3):\n        super().__init__()\n        self.save_hyperparameters()\n        self.continuous_scheduler = continuous_scheduler\n        # Layers\n        self.feature_extractor = timm.create_model(CFG.MODEL_NAME, pretrained=True)\n        in_features = self.feature_extractor.classifier.in_features\n        self.feature_extractor.classifier = nn.Identity()\n        self.feature_extractor.global_pool = nn.Identity()\n        self.pooling = GeM()\n        self.dropout = nn.Dropout()\n        self.dense = nn.Linear(in_features, CFG.EMBEDDING_SIZE)\n        self.swish = Swish_module()\n        self.metric_classify = ArcMarginProduct_subcenter(CFG.EMBEDDING_SIZE)\n        # Loss\n        self.criterion = ArcFaceLossAdaptiveMargin(margins=MARGINS)\n        # Metrics\n        self.train_acc = torchmetrics.Accuracy()\n        self.train_top_k_acc = torchmetrics.Accuracy(top_k=5)\n        self.val_acc = torchmetrics.Accuracy()\n        self.val_top_k_acc = torchmetrics.Accuracy(top_k=5)\n        \n    def forward(self, image):\n        \"\"\"\n            Return embedding of the images\n        \"\"\"\n        features = self.feature_extractor(image)\n        x = self.pooling(features).flatten(1)\n        x = self.dropout(x)\n        x = self.dense(x)\n        return F.normalize(x)\n    \n    def training_step(self, batch, batch_idx):\n        \"\"\"\n            Return the loss to do a step on\n        \"\"\"\n        img, label = batch\n        embedding = self(img)\n        logits = self.metric_classify(embedding)\n        loss  = self.criterion(logits, label)\n        # Log metrics\n        self.train_acc(logits, label)\n        self.train_top_k_acc(logits, label)\n        self.log(\"train/loss\", loss)\n        self.log(\"train/acc\", self.train_acc)\n        self.log(\"train/top_k_acc\", self.train_top_k_acc)\n        # Return loss, labels and preds\n        return {\"loss\": loss, \"preds\": logits.detach(), \"targets\": label.detach()}\n    \n    def configure_optimizers(self):\n        \"\"\"\n            Build optimizer(s) and lr scheduler(s)\n        \"\"\"\n        optimizer = torch.optim.AdamW(self.parameters())\n        if self.continuous_scheduler:\n            sched = {\n                \"scheduler\": PeakScheduler(optimizer, epoch_size=len(train_loader) // CFG.ACCUMULATE_GRAD_BATCHES, verbose=False),\n                \"interval\": \"step\",\n            }\n        else:\n            sched = {\n                \"scheduler\": PeakScheduler(optimizer),\n                \"interval\": \"epoch\",\n            }\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": sched\n        }\n    \n    def save_class_weights(self):\n        \"\"\"\n            Save the class centers as a tensor\n        \"\"\"\n        torch.save(self.metric_classify.weight, 'class_weights.pt')","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:32.032081Z","iopub.execute_input":"2022-03-07T18:25:32.032490Z","iopub.status.idle":"2022-03-07T18:25:32.055906Z","shell.execute_reply.started":"2022-03-07T18:25:32.032454Z","shell.execute_reply":"2022-03-07T18:25:32.055165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training 🏃\nCreate a Pytorch lightning trainer with our configuration, and run our model on our dataset:","metadata":{}},{"cell_type":"code","source":"model = WandDIDNet()\nwandb_logger = WandbLogger(project=\"W&D - identification\")\n# Trainer \ntrainer = Trainer(\n    profiler=\"simple\", # Profiling\n    accumulate_grad_batches=CFG.ACCUMULATE_GRAD_BATCHES,# Accumulate gradient over multiple batches\n    gpus=1,# Use the one GPU we have\n    precision=16 if CFG.MIXED_PRECISION else 32,# Mixed precision\n    max_epochs=CFG.EPOCHS,\n    logger=wandb_logger,\n    log_every_n_steps=10\n)\n# Let's go ⚡\ntrainer.fit(model, train_loader)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:25:32.057455Z","iopub.execute_input":"2022-03-07T18:25:32.057897Z","iopub.status.idle":"2022-03-07T18:26:11.630863Z","shell.execute_reply.started":"2022-03-07T18:25:32.057862Z","shell.execute_reply":"2022-03-07T18:26:11.629759Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Save everything\ntrainer.save_checkpoint(CFG.MODEL_PATH)\nartifact = wandb.log_artifact(CFG.MODEL_PATH, name='w_and_d-id-normal', type='model') \nmodel.save_class_weights()\nartifact = wandb.log_artifact('class_weights.pt', name='w_and_d-id-normal-weights', type=\"class_weights\") ","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:26:11.636182Z","iopub.execute_input":"2022-03-07T18:26:11.636514Z","iopub.status.idle":"2022-03-07T18:26:17.312903Z","shell.execute_reply.started":"2022-03-07T18:26:11.636447Z","shell.execute_reply":"2022-03-07T18:26:17.312178Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference 🔮\nNow we have to do the inference part. First, we will predict embeddings for the training and testing samples. Then we will use nearest neighboors to make our predictions","metadata":{}},{"cell_type":"code","source":"class WandDIDPred(Dataset):\n    def __init__(self, data, folder, base_path=BASE_PATH):\n        self.base_path = os.path.join(base_path, folder)\n        self.data = data\n        # Augmentations\n        transformations = albumentations.Compose([\n            albumentations.Normalize(),\n            ToTensorV2(p=1.0)\n        ])\n\n        def make_transform(transform=False):\n            def f(image):\n                if transform:\n                    image_np = np.array(image)\n                    augmented = transform(image=image_np)\n                return augmented\n            return f\n\n        self.transforms = transforms.Compose([\n            transforms.Lambda(make_transform(transformations)),\n        ])\n        \n    def __getitem__(self, idx):\n        image = self.preprocess(self.data[\"image\"].iloc[idx])\n        return image\n    \n    def preprocess(self, image):\n        image = os.path.join(self.base_path, image)\n        image = cv2.imread(image)[:, :, ::-1]\n        if image.shape[0] < CFG.IMAGE_SIZE or image.shape[1] < CFG.IMAGE_SIZE:\n            image = cv2.resize(image, (CFG.IMAGE_SIZE, CFG.IMAGE_SIZE), cv2.INTER_CUBIC)\n        if self.transforms is not None:\n            image = self.transforms(image)[\"image\"]\n        return image\n    \n    def __len__(self):\n        return len(self.data)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:26:17.314282Z","iopub.execute_input":"2022-03-07T18:26:17.314531Z","iopub.status.idle":"2022-03-07T18:26:17.326602Z","shell.execute_reply.started":"2022-03-07T18:26:17.314497Z","shell.execute_reply":"2022-03-07T18:26:17.325502Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Prediction on the training data\npred_loader = torch.utils.data.DataLoader(\n    WandDIDPred(data, \"train_images\"),\n    batch_size=CFG.BATCH_SIZE * 6,\n    num_workers=CFG.NUM_WORKERS,\n    pin_memory=True,\n    shuffle=False\n)\npreds = trainer.predict(model, dataloaders=pred_loader)\npreds = torch.cat(preds, dim=0)\ntrain_data = data.copy()\ntrain_data[\"embedding\"] = preds.tolist()\ntrain_data.to_csv(\"train.csv\")\ntrain_data","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:26:17.328133Z","iopub.execute_input":"2022-03-07T18:26:17.328640Z","iopub.status.idle":"2022-03-07T18:40:15.866342Z","shell.execute_reply.started":"2022-03-07T18:26:17.328602Z","shell.execute_reply":"2022-03-07T18:40:15.865542Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Prediction on test data\ntest_data = pd.read_csv(os.path.join(BASE_PATH, \"sample_submission.csv\"), index_col=\"image\")\nif \"inference_image\" in test_data.columns:\n    test_data[\"image\"] = test_data[\"inference_image\"]\nelse:\n    test_data[\"image\"] = test_data.index\n    \npred_loader = torch.utils.data.DataLoader(\n    WandDIDPred(test_data, \"test_images\"),\n    batch_size=CFG.BATCH_SIZE * 6,\n    num_workers=CFG.NUM_WORKERS,\n    pin_memory=True,\n    shuffle=False\n)\npreds = trainer.predict(model, dataloaders=pred_loader)\npreds = torch.cat(preds, dim=0)\ntest_data[\"embedding\"] = preds.tolist()\ntest_data.to_csv(\"test.csv\")\ntest_data","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:44:09.684932Z","iopub.execute_input":"2022-03-07T18:44:09.685505Z","iopub.status.idle":"2022-03-07T18:44:25.284257Z","shell.execute_reply.started":"2022-03-07T18:44:09.685457Z","shell.execute_reply":"2022-03-07T18:44:25.282200Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Log our model in wandb and finish the run\nwandb_logger.finalize(\"success\")\nwandb.finish()","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:40:16.638871Z","iopub.status.idle":"2022-03-07T18:40:16.639830Z","shell.execute_reply.started":"2022-03-07T18:40:16.639528Z","shell.execute_reply":"2022-03-07T18:40:16.639559Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_data[\"individual_id_integer\"] = train_data[\"individual_id\"].map(individual_mapping).fillna(-1)\ntrain_embeddings = np.array(train_data[\"embedding\"].values.tolist())\ntest_embeddings = np.array(test_data[\"embedding\"].values.tolist())\nclass_centers = model.metric_classify.weight.detach().numpy()","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:40:16.640857Z","iopub.status.idle":"2022-03-07T18:40:16.641658Z","shell.execute_reply.started":"2022-03-07T18:40:16.641403Z","shell.execute_reply":"2022-03-07T18:40:16.641433Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Solution:\n    def __init__(self, database, q_prior):\n        self.database_embeddings = np.array(database[\"embeddings\"][\"embedding\"].values.tolist())\n        self.database_individuals = database[\"embeddings\"][\"individual_id\"].values\n        self.q_prior = q_prior\n        self.embed_neigh = NearestNeighbors(n_neighbors=CFG.KNN,metric='cosine')\n        self.embed_neigh.fit(self.database_embeddings)\n        self.class_neigh = NearestNeighbors(n_neighbors=CFG.KNN,metric='cosine')\n        self.class_neigh.fit(database[\"class_centers\"])\n        self.default = ['938b7e931166', '5bf17305f073', '7593d2aee842', '7362d7a01d00','956562ff2888']        \n    \n    def predict(self, queries):\n        embed_distances, embed_idxs = self.embed_neigh.kneighbors(queries, CFG.KNN, return_distance=True)\n        class_distances, class_idxs = self.class_neigh.kneighbors(queries, CFG.KNN, return_distance=True)\n        \n        class_individuals = np.repeat(list(individual_mapping), CFG.CENTERS_PER_CLASS)[class_idxs]\n        embed_individuals = self.database_individuals[embed_idxs]\n        \n        n = embed_distances.size\n        embeddings_df = pd.DataFrame(data={\n            'distance': embed_distances.ravel(),\n            'individual': embed_individuals.ravel(),\n            'query_id': np.repeat(np.arange(len(queries)), CFG.KNN)\n        }, index=np.arange(n))\n        \n        class_df = pd.DataFrame(data={\n            'distance': class_distances.ravel(),\n            'individual': class_individuals.ravel(),\n            'query_id': np.repeat(np.arange(len(queries)), CFG.KNN)\n        }, index=np.arange(n))\n        \n        embeddings_topk = embeddings_df.groupby([\"query_id\", \"individual\"]).agg(\"min\")['distance'].groupby('query_id', group_keys=False).nsmallest(5)\n        class_topk = class_df.groupby([\"query_id\", \"individual\"]).agg(\"min\")['distance'].groupby('query_id', group_keys=False).nsmallest(5)\n        embeddings_topk = embeddings_topk.reset_index().groupby(\"query_id\").agg(list)\n        class_topk = class_topk.reset_index().groupby(\"query_id\").agg(list)\n        class_t_new = np.quantile(class_topk[\"distance\"].apply(lambda x: x[0]), 1 - self.q_prior)\n        embeddings_t_new = np.quantile(embeddings_topk[\"distance\"].apply(lambda x: x[0]), 1 - self.q_prior)\n        \n        def insert_new_individuals(x):\n            m = np.array(x[\"distance\"]) > class_t_new\n            preds = x[\"individual\"]\n            if m.any():\n                preds.insert(np.argmax(m), \"new_individual\")\n            preds = preds + [y for y in self.default if y not in preds]\n            return preds[:5]\n        \n        preds = class_topk.apply(insert_new_individuals, axis=1)\n        return preds.values.tolist()","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:40:16.642991Z","iopub.status.idle":"2022-03-07T18:40:16.643724Z","shell.execute_reply.started":"2022-03-07T18:40:16.643459Z","shell.execute_reply":"2022-03-07T18:40:16.643490Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<p style=\"background: #ff6363; padding: 25px; border-radius: 10px\">\n    🚨 The current solution does not utilize image samples (uses only class centers) so their is room for improvement here ! 🚨\n</p>","metadata":{}},{"cell_type":"code","source":"solution = Solution({\n    \"embeddings\": train_data,\n    \"class_centers\": class_centers\n}, CFG.Q_NEW)\npredictions = solution.predict(test_embeddings)\npredictions = pd.Series(predictions, test_data.index, name=\"predictions\").map(lambda x: \" \".join(x))\npredictions.to_csv(\"submission.csv\")\npredictions","metadata":{"execution":{"iopub.status.busy":"2022-03-07T18:40:16.644925Z","iopub.status.idle":"2022-03-07T18:40:16.645666Z","shell.execute_reply.started":"2022-03-07T18:40:16.645427Z","shell.execute_reply":"2022-03-07T18:40:16.645454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"![](https://upload.wikimedia.org/wikipedia/commons/thumb/e/ea/Thats_all_folks.svg/2560px-Thats_all_folks.svg.png)","metadata":{}}]}