{"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":"# 🐋🐬⚡ PyTorch Lightning ArcFace Train & Infer\n\nLet's train [`timm`](https://github.com/rwightman/pytorch-image-models) models \nwith [PyTorch Lightning](https://www.pytorchlightning.ai/)!\n\n**Sources:**\n- [[Pytorch] ArcFace + GeM Pooling Starter](https://www.kaggle.com/debarshichanda/pytorch-arcface-gem-pooling-starter)\n- [FAISS Pytorch Inference](https://www.kaggle.com/debarshichanda/faiss-pytorch-inference)\n- [whale2-cropped-dataset](https://www.kaggle.com/phalanx/whale2-cropped-dataset)\n\n**Nice Lightning `Trainer` Flags to try:**\n- Accumulate gradient every k batches: `accumulate_grad_batches=2`\n- Run a learning rate finder algorithm: `auto_lr_find=True`\n- Automatically try to find the largest batch size that fits into memory: `auto_scale_batch_size=True`\n- Quickly check whether everything runs fine: `fast_dev_run=True`\n- Train on multiple GPUs: `gpus=2` (if you use multiple GPUs, also set `accelerator=ddp`)\n- Train with half precision: `precision=16`\n- Use Stochastic Weight Averaging: `stochastic_weight_avg=True`\n\n**Public LB scores**\n- V02: 0.378 (`image_size=256`, `\"tf_efficientnet_b2\"`, `batch_size=128`, `learning_rate=3e-4`)\n- V10: 0.439 (`image_size=512`, `\"tf_efficientnet_b0\"`, `batch_size=64`, `learning_rate=3e-4`)\n- V12: 0.498 (`image_size=512`, `\"tf_efficientnet_b2\"`, `batch_size=32`, `learning_rate=3e-4`)\n- V14: 0.567 (`image_size=512`, `\"tf_efficientnet_b4\"`, `batch_size=16`, `learning_rate=3e-4`)\n- V21: 0.656 (backfin cropped data, `image_size=384`, `\"tf_efficientnet_b4\"`, `batch_size=32`, `learning_rate=3e-4`, `scheduler=OneCycleLR`)\n- V22: 0.XXX (backfin cropped data, `image_size=384`, `\"convnext_small\"`, `batch_size=32`, `learning_rate=3e-4`, `scheduler=OneCycleLR`)","metadata":{}},{"cell_type":"markdown","source":"# Installations (timm + FAISS)","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.append(\"../input/timmmaster\")\nsys.path.append(\"\")\n!pip install ../input/faissgpu17/faiss_gpu-1.7.0-cp37-cp37m-manylinux2014_x86_64.whl","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:19.570726Z","iopub.execute_input":"2022-04-16T01:09:19.570986Z","iopub.status.idle":"2022-04-16T01:09:48.623815Z","shell.execute_reply.started":"2022-04-16T01:09:19.570955Z","shell.execute_reply":"2022-04-16T01:09:48.622937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import math\nfrom typing import Callable\nfrom typing import Dict\nfrom typing import Optional\nfrom typing import Tuple\nfrom pathlib import Path\n\nimport faiss\nimport numpy as np\nimport pandas as pd\nimport pytorch_lightning as pl\nimport timm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom PIL import Image\nfrom pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint\nfrom timm.data.transforms_factory import create_transform\nfrom timm.optim import create_optimizer_v2\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data import Dataset\nfrom tqdm.notebook import tqdm\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.preprocessing import normalize\nfrom sklearn.preprocessing import LabelEncoder","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-04-16T01:09:48.626089Z","iopub.execute_input":"2022-04-16T01:09:48.626389Z","iopub.status.idle":"2022-04-16T01:09:48.635037Z","shell.execute_reply.started":"2022-04-16T01:09:48.626353Z","shell.execute_reply":"2022-04-16T01:09:48.634302Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Paths & Settings","metadata":{}},{"cell_type":"code","source":"INPUT_DIR = Path(\"..\") / \"input\"\nOUTPUT_DIR = Path(\"/\") / \"kaggle\" / \"working\"\n\nDATA_ROOT_DIR = INPUT_DIR / \"convert-backfintfrecords\" / \"happy-whale-and-dolphin-backfin\"\nTRAIN_DIR = DATA_ROOT_DIR / \"train_images\"\nTEST_DIR = DATA_ROOT_DIR / \"test_images\"\nTRAIN_CSV_PATH = DATA_ROOT_DIR / \"train.csv\"\nSAMPLE_SUBMISSION_CSV_PATH = DATA_ROOT_DIR /\"sample_submission.csv\"\nPUBLIC_SUBMISSION_CSV_PATH = INPUT_DIR / \"publicsubmission758\" / \"submission.csv\"\nIDS_WITHOUT_BACKFIN_PATH = INPUT_DIR / \"ids-without-backfin\" / \"ids_without_backfin.npy\"\n\nN_SPLITS = 5\n\nENCODER_CLASSES_PATH = OUTPUT_DIR /\" encoder_classes.npy\"\nTEST_CSV_PATH = OUTPUT_DIR / \" test.csv\"\nTRAIN_CSV_ENCODED_FOLDED_PATH = OUTPUT_DIR / \"train_encoded_folded.csv\"\nCHECKPOINTS_DIR = \"../input/conv-nextday322pl\"\nSUBMISSION_CSV_PATH = OUTPUT_DIR / \"submission.csv\"\n\nDEBUG = False","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:48.636171Z","iopub.execute_input":"2022-04-16T01:09:48.636935Z","iopub.status.idle":"2022-04-16T01:09:48.645794Z","shell.execute_reply.started":"2022-04-16T01:09:48.636897Z","shell.execute_reply":"2022-04-16T01:09:48.645014Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prepare DataFrames","metadata":{}},{"cell_type":"code","source":"def get_image_path(id: str, dir: Path) -> str:\n    return f\"{dir / id}\"","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:48.649539Z","iopub.execute_input":"2022-04-16T01:09:48.651512Z","iopub.status.idle":"2022-04-16T01:09:48.660163Z","shell.execute_reply.started":"2022-04-16T01:09:48.651431Z","shell.execute_reply":"2022-04-16T01:09:48.659291Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train DataFrame","metadata":{}},{"cell_type":"code","source":"train_df = pd.read_csv(TRAIN_CSV_PATH)\n\ntrain_df[\"image_path\"] = train_df[\"image\"].apply(get_image_path, dir=TRAIN_DIR)\n\nencoder = LabelEncoder()\ntrain_df[\"individual_id\"] = encoder.fit_transform(train_df[\"individual_id\"])\nnp.save(ENCODER_CLASSES_PATH, encoder.classes_)\n\nskf = StratifiedKFold(n_splits=N_SPLITS)\nfor fold, (_, val_) in enumerate(skf.split(X=train_df, y=train_df.individual_id)):\n    train_df.loc[val_, \"kfold\"] = fold\n    \ntrain_df.to_csv(TRAIN_CSV_ENCODED_FOLDED_PATH, index=False)\n    \ntrain_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:48.661562Z","iopub.execute_input":"2022-04-16T01:09:48.662145Z","iopub.status.idle":"2022-04-16T01:09:49.713925Z","shell.execute_reply.started":"2022-04-16T01:09:48.662099Z","shell.execute_reply":"2022-04-16T01:09:49.713200Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Test DataFrame","metadata":{}},{"cell_type":"code","source":"# Use sample submission csv as template\ntest_df = pd.read_csv(SAMPLE_SUBMISSION_CSV_PATH)\ntest_df[\"image_path\"] = test_df[\"image\"].apply(get_image_path, dir=TEST_DIR)\n\ntest_df.drop(columns=[\"predictions\"], inplace=True)\n\n# Dummy id\ntest_df[\"individual_id\"] = 0\n\ntest_df.to_csv(TEST_CSV_PATH, index=False)\n\ntest_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:49.716330Z","iopub.execute_input":"2022-04-16T01:09:49.716739Z","iopub.status.idle":"2022-04-16T01:09:50.092711Z","shell.execute_reply.started":"2022-04-16T01:09:49.716700Z","shell.execute_reply":"2022-04-16T01:09:50.092000Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class HappyWhaleDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, transform: Optional[Callable] = None):\n        self.df = df\n        self.transform = transform\n\n        self.image_names = self.df[\"image\"].values\n        self.image_paths = self.df[\"image_path\"].values\n        self.targets = self.df[\"individual_id\"].values\n\n    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:\n        image_name = self.image_names[index]\n\n        image_path = self.image_paths[index]\n\n        image = Image.open(image_path).convert(\"RGB\")\n        \n        if self.transform:\n            image = self.transform(image)\n\n        target = self.targets[index]\n        target = torch.tensor(target, dtype=torch.long)\n\n        return {\"image_name\": image_name, \"image\": image, \"target\": target}\n\n    def __len__(self) -> int:\n        return len(self.df)","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:50.093879Z","iopub.execute_input":"2022-04-16T01:09:50.095437Z","iopub.status.idle":"2022-04-16T01:09:50.104062Z","shell.execute_reply.started":"2022-04-16T01:09:50.095386Z","shell.execute_reply":"2022-04-16T01:09:50.103304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Lightning DataModule","metadata":{}},{"cell_type":"code","source":"class LitDataModule(pl.LightningDataModule):\n    def __init__(\n        self,\n        train_csv_encoded_folded: str,\n        test_csv: str,\n        val_fold: float,\n        image_size: int,\n        batch_size: int,\n        num_workers: int,\n    ):\n        super().__init__()\n\n        self.save_hyperparameters()\n\n        self.train_df = pd.read_csv(train_csv_encoded_folded)\n        self.test_df = pd.read_csv(test_csv)\n        \n        self.transform = create_transform(\n            input_size=(self.hparams.image_size, self.hparams.image_size),\n            crop_pct=1.0,\n        )\n        \n    def setup(self, stage: Optional[str] = None):\n        if stage == \"fit\" or stage is None:\n            # Split train df using fold\n            train_df = self.train_df[self.train_df.kfold != self.hparams.val_fold].reset_index(drop=True)\n            val_df = self.train_df[self.train_df.kfold == self.hparams.val_fold].reset_index(drop=True)\n\n            self.train_dataset = HappyWhaleDataset(train_df, transform=self.transform)\n            self.val_dataset = HappyWhaleDataset(val_df, transform=self.transform)\n\n        if stage == \"test\" or stage is None:\n            self.test_dataset = HappyWhaleDataset(self.test_df, transform=self.transform)\n\n    def train_dataloader(self) -> DataLoader:\n        return self._dataloader(self.train_dataset, train=True)\n\n    def val_dataloader(self) -> DataLoader:\n        return self._dataloader(self.val_dataset)\n\n    def test_dataloader(self) -> DataLoader:\n        return self._dataloader(self.test_dataset)\n\n    def _dataloader(self, dataset: HappyWhaleDataset, train: bool = False) -> DataLoader:\n        return DataLoader(\n            dataset,\n            batch_size=self.hparams.batch_size,\n            shuffle=train,\n            num_workers=self.hparams.num_workers,\n            pin_memory=True,\n            drop_last=train,\n        )","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:50.105931Z","iopub.execute_input":"2022-04-16T01:09:50.106143Z","iopub.status.idle":"2022-04-16T01:09:50.121839Z","shell.execute_reply.started":"2022-04-16T01:09:50.106119Z","shell.execute_reply":"2022-04-16T01:09:50.121017Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ArcMargin","metadata":{}},{"cell_type":"code","source":"class CrossEntropyLabelSmooth(nn.Module):\n    \"\"\"Cross entropy loss with label smoothing regularizer.\n    Reference:\n    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.\n    Equation: y = (1 - epsilon) * y + epsilon / K.\n    Args:\n        num_classes (int): number of classes.\n        epsilon (float): weight.\n    \"\"\"\n\n    def __init__(self, num_classes, epsilon=0.1, use_gpu=True):\n        super(CrossEntropyLabelSmooth, self).__init__()\n        self.num_classes = num_classes\n        self.epsilon = epsilon\n        self.use_gpu = use_gpu\n        self.logsoftmax = nn.LogSoftmax(dim=1)\n\n    def forward(self, inputs, targets):\n        \"\"\"\n        Args:\n            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)\n            targets: ground truth labels with shape (num_classes)\n        \"\"\"\n        log_probs = self.logsoftmax(inputs)\n        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1)\n        if self.use_gpu: targets = targets.cuda()\n        targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes\n        loss = (- targets * log_probs).mean(0).sum()\n        return loss\n\n\n\n# From https://github.com/lyakaap/Landmark2019-1st-and-3rd-Place-Solution/blob/master/src/modeling/metric_learning.py\n# Added type annotations, device, and 16bit support\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\n    def __init__(\n        self,\n        in_features: int,\n        out_features: int,\n        s: float,\n        m: float,\n        easy_margin: bool,\n        ls_eps: float,\n    ):\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.ls_eps = ls_eps  # label smoothing\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: torch.Tensor, label: torch.Tensor, device: str = \"cuda\") -> torch.Tensor:\n        # --------------------------- cos(theta) & phi(theta) ---------------------\n        cosine = F.linear(F.normalize(input), F.normalize(self.weight))\n        # Enable 16 bit precision\n        cosine = cosine.to(torch.float32)\n\n        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))\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=device)\n        one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n        if self.ls_eps > 0:\n            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features\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","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:50.124952Z","iopub.execute_input":"2022-04-16T01:09:50.125211Z","iopub.status.idle":"2022-04-16T01:09:50.145334Z","shell.execute_reply.started":"2022-04-16T01:09:50.125179Z","shell.execute_reply":"2022-04-16T01:09:50.144310Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Lightning Module","metadata":{}},{"cell_type":"code","source":"def l2_norm(input, axis=1):\n    norm = torch.norm(input, 2, axis, True)\n    output = torch.div(input, norm)\n    return output\nfrom torch.autograd import Variable\n\n\nclass FocalLoss(nn.Module):\n    def __init__(self, gamma=0, alpha=None, size_average=True):\n        super(FocalLoss, self).__init__()\n        self.gamma = gamma\n        self.alpha = alpha\n        if isinstance(alpha, (float, int)): self.alpha = torch.Tensor([alpha, 1 - alpha])\n        if isinstance(alpha, list): self.alpha = torch.Tensor(alpha)\n        self.size_average = size_average\n\n    def forward(self, input, target):\n        if input.dim() > 2:\n            input = input.view(input.size(0), input.size(1), -1)  # N,C,H,W => N,C,H*W\n            input = input.transpose(1, 2)  # N,C,H*W => N,H*W,C\n            input = input.contiguous().view(-1, input.size(2))  # N,H*W,C => N*H*W,C\n        target = target.view(-1, 1)\n\n        logpt = F.log_softmax(input)\n        logpt = logpt.gather(1, target)\n        logpt = logpt.view(-1)\n        pt = Variable(logpt.data.exp())\n\n        if self.alpha is not None:\n            if self.alpha.type() != input.data.type():\n                self.alpha = self.alpha.type_as(input.data)\n            select = (target != 0).type(torch.LongTensor).cuda()\n            at = self.alpha.gather(0, select.data.view(-1))\n            logpt = logpt * Variable(at)\n\n        loss = -1 * (1 - pt) ** self.gamma * logpt\n        if self.size_average:\n            return loss.mean()\n        else:\n            return loss.sum()\n\nclass LitModule(pl.LightningModule):\n    def __init__(\n            self,\n            model_name: str,\n            pretrained: bool,\n            drop_rate: float,\n            embedding_size: int,\n            num_classes: int,\n            arc_s: float,\n            arc_m: float,\n            arc_easy_margin: bool,\n            arc_ls_eps: float,\n            optimizer: str,\n            learning_rate: float,\n            weight_decay: float,\n            len_train_dl: int,\n            epochs: int\n    ):\n        super().__init__()\n\n        self.save_hyperparameters()\n        # self.fea_extra_layer = [2, 3]\n        self.fea_extra_layer = [-2, -1]\n        self.model = timm.create_model(model_name, pretrained=False, drop_rate=drop_rate,\n                                       features_only=True,\n                                       out_indices=self.fea_extra_layer\n                                       )\n        in_features = 1536\n        self.embedding = nn.Sequential(\n\n            nn.Linear(in_features, embedding_size),\n            nn.BatchNorm1d(embedding_size))\n\n        # self.model.reset_classifier(num_classes=0, global_pool=\"avg\")\n        self.bn = nn.Sequential(\n            nn.BatchNorm2d(1024),\n            nn.Dropout(0.2),\n            nn.AdaptiveAvgPool2d(1),\n        )\n        self.bn2 = nn.Sequential(\n            nn.BatchNorm2d(512),\n            nn.Dropout(0.2),\n            nn.AdaptiveAvgPool2d(1),\n        )\n\n        self.arc = ArcMarginProduct(\n            in_features=embedding_size,\n            out_features=num_classes,\n            s=arc_s,\n            m=arc_m,\n            easy_margin=arc_easy_margin,\n            ls_eps=arc_ls_eps,\n        )\n        self.loss_fn = nn.CrossEntropyLoss()\n\n    def forward(self, images: torch.Tensor) -> torch.Tensor:\n        features = self.model(images)\n        features[0] = self.bn2(features[0])\n        features[1] = self.bn(features[1])\n        features = torch.cat(features, dim=1)\n        embeddings = self.embedding(features.flatten(1))\n\n        return embeddings\n\n    def configure_optimizers(self):\n        optimizer = create_optimizer_v2(\n            self.parameters(),\n            opt=self.hparams.optimizer,\n            lr=self.hparams.learning_rate,\n            weight_decay=self.hparams.weight_decay,\n        )\n\n        scheduler = torch.optim.lr_scheduler.OneCycleLR(\n            optimizer,\n            self.hparams.learning_rate,\n            steps_per_epoch=self.hparams.len_train_dl,\n            epochs=self.hparams.epochs,\n        )\n        scheduler = {\"scheduler\": scheduler, \"interval\": \"step\"}\n\n        return [optimizer], [scheduler]\n\n    def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor:\n        return self._step(batch, \"train\")\n\n    def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor:\n        return self._step(batch, \"val\")\n\n    def _step(self, batch: Dict[str, torch.Tensor], step: str) -> torch.Tensor:\n        images, targets = batch[\"image\"], batch[\"target\"]\n\n        embeddings = self(images)\n        outputs = self.arc(embeddings, targets, self.device)\n\n        loss = self.loss_fn(outputs, targets)\n\n        self.log(f\"{step}_loss\", loss)\n\n        return loss","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:50.149129Z","iopub.execute_input":"2022-04-16T01:09:50.149369Z","iopub.status.idle":"2022-04-16T01:09:50.178159Z","shell.execute_reply.started":"2022-04-16T01:09:50.149340Z","shell.execute_reply":"2022-04-16T01:09:50.177444Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"def train(\n    train_csv_encoded_folded: str = str(TRAIN_CSV_ENCODED_FOLDED_PATH),\n    test_csv: str = str(TEST_CSV_PATH),\n    val_fold: float = 0.0,\n    image_size: int = 256,\n    batch_size: int = 64,\n    num_workers: int = 2,\n    model_name: str = \"tf_efficientnet_b0\",\n    pretrained: bool = False,\n    drop_rate: float = 0.0,\n    embedding_size: int = 512,\n    num_classes: int = 15587,\n    arc_s: float = 64.0,\n    arc_m: float = 0.5,\n    arc_easy_margin: bool = False,\n    arc_ls_eps: float = 0.0,\n    optimizer: str = \"adam\",\n    learning_rate: float = 3e-4,\n    weight_decay: float = 1e-6,\n    checkpoints_dir: str = str(CHECKPOINTS_DIR),\n    accumulate_grad_batches: int = 1,\n    auto_lr_find: bool = False,\n    auto_scale_batch_size: bool = False,\n    fast_dev_run: bool = False,\n    gpus: int = 1,\n    max_epochs: int = 0,\n    precision: int = 16,\n    stochastic_weight_avg: bool = True,\n):\n    pl.seed_everything(42)\n\n    datamodule = LitDataModule(\n        train_csv_encoded_folded=train_csv_encoded_folded,\n        test_csv=test_csv,\n        val_fold=val_fold,\n        image_size=image_size,\n        batch_size=batch_size,\n        num_workers=num_workers,\n    )\n    \n    datamodule.setup()\n    len_train_dl = len(datamodule.train_dataloader())\n\n    module = LitModule(\n        model_name=model_name,\n        pretrained=pretrained,\n        drop_rate=drop_rate,\n        embedding_size=embedding_size,\n        num_classes=num_classes,\n        arc_s=arc_s,\n        arc_m=arc_m,\n        arc_easy_margin=arc_easy_margin,\n        arc_ls_eps=arc_ls_eps,\n        optimizer=optimizer,\n        learning_rate=learning_rate,\n        weight_decay=weight_decay,\n        len_train_dl=len_train_dl,\n        epochs=max_epochs\n    )\n    \n    model_checkpoint = ModelCheckpoint(\n        checkpoints_dir,\n        filename=f\"{model_name}_{image_size}\",\n        monitor=\"val_loss\",\n    )\n        \n    trainer = pl.Trainer(\n        accumulate_grad_batches=accumulate_grad_batches,\n        auto_lr_find=auto_lr_find,\n        auto_scale_batch_size=auto_scale_batch_size,\n        benchmark=True,\n        callbacks=[model_checkpoint],\n        deterministic=True,\n        fast_dev_run=fast_dev_run,\n        gpus=gpus,\n        max_epochs=2 if DEBUG else max_epochs,\n        precision=precision,\n        stochastic_weight_avg=stochastic_weight_avg,\n        limit_train_batches=0.1 if DEBUG else 1.0,\n        limit_val_batches=0.1 if DEBUG else 1.0,\n    )\n\n    trainer.tune(module, datamodule=datamodule)\n\n    trainer.fit(module, datamodule=datamodule)","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:50.181337Z","iopub.execute_input":"2022-04-16T01:09:50.181722Z","iopub.status.idle":"2022-04-16T01:09:50.196908Z","shell.execute_reply.started":"2022-04-16T01:09:50.181688Z","shell.execute_reply":"2022-04-16T01:09:50.196208Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"def load_eval_module(checkpoint_path: str, device: torch.device) -> LitModule:\n    module = LitModule.load_from_checkpoint(checkpoint_path)\n    module.to(device)\n    module.eval()\n\n    return module\n\ndef load_dataloaders(\n    train_csv_encoded_folded: str,\n    test_csv: str,\n    val_fold: float,\n    image_size: int,\n    batch_size: int,\n    num_workers: int,\n) -> Tuple[DataLoader, DataLoader, DataLoader]:\n\n    datamodule = LitDataModule(\n        train_csv_encoded_folded=train_csv_encoded_folded,\n        test_csv=test_csv,\n        val_fold=val_fold,\n        image_size=image_size,\n        batch_size=batch_size,\n        num_workers=num_workers,\n        \n    )\n\n    datamodule.setup()\n\n    train_dl = datamodule.train_dataloader()\n    val_dl = datamodule.val_dataloader()\n    test_dl = datamodule.test_dataloader()\n\n    return train_dl, val_dl, test_dl\n\n\ndef load_encoder() -> LabelEncoder:\n    encoder = LabelEncoder()\n    encoder.classes_ = np.load(ENCODER_CLASSES_PATH, allow_pickle=True)\n\n    return encoder\n\n\n@torch.inference_mode()\ndef get_embeddings(\n    module: pl.LightningModule, dataloader: DataLoader, encoder: LabelEncoder, stage: str\n) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n\n    all_image_names = []\n    all_embeddings = []\n    all_targets = []\n\n    for batch in tqdm(dataloader, desc=f\"Creating {stage} embeddings\"):\n        image_names = batch[\"image_name\"]\n        images = batch[\"image\"].to(module.device)\n        targets = batch[\"target\"].to(module.device)\n\n        embeddings = module(images)\n        #emb2 = module2(images)\n        #embeddings = embeddings + emb2   \n        all_image_names.append(image_names)\n        all_embeddings.append(embeddings.cpu().numpy())\n        all_targets.append(targets.cpu().numpy())\n        \n        if DEBUG:\n            break\n\n    all_image_names = np.concatenate(all_image_names)\n    all_embeddings = np.vstack(all_embeddings)\n    all_targets = np.concatenate(all_targets)\n\n    all_embeddings = normalize(all_embeddings, axis=1, norm=\"l2\")\n    all_targets = encoder.inverse_transform(all_targets)\n\n    return all_image_names, all_embeddings, all_targets\n\n\ndef create_and_search_index(embedding_size: int, train_embeddings: np.ndarray, val_embeddings: np.ndarray, k: int):\n    index = faiss.IndexFlatIP(embedding_size)\n    index.add(train_embeddings)\n    D, I = index.search(val_embeddings, k=k)  # noqa: E741\n\n    return D, I\n\n\ndef create_val_targets_df(\n    train_targets: np.ndarray, val_image_names: np.ndarray, val_targets: np.ndarray\n) -> pd.DataFrame:\n\n    allowed_targets = np.unique(train_targets)\n    val_targets_df = pd.DataFrame(np.stack([val_image_names, val_targets], axis=1), columns=[\"image\", \"target\"])\n    val_targets_df.loc[~val_targets_df.target.isin(allowed_targets), \"target\"] = \"new_individual\"\n\n    return val_targets_df\n\n\ndef create_distances_df(\n    image_names: np.ndarray, targets: np.ndarray, D: np.ndarray, I: np.ndarray, stage: str  # noqa: E741\n) -> pd.DataFrame:\n\n    distances_df = []\n    for i, image_name in tqdm(enumerate(image_names), desc=f\"Creating {stage}_df\"):\n        target = targets[I[i]]\n        distances = D[i]\n        subset_preds = pd.DataFrame(np.stack([target, distances], axis=1), columns=[\"target\", \"distances\"])\n        subset_preds[\"image\"] = image_name\n        distances_df.append(subset_preds)\n\n    distances_df = pd.concat(distances_df).reset_index(drop=True)\n    distances_df = distances_df.groupby([\"image\", \"target\"]).distances.max().reset_index()\n    distances_df = distances_df.sort_values(\"distances\", ascending=False).reset_index(drop=True)\n\n    return distances_df\n\n\ndef get_best_threshold(val_targets_df: pd.DataFrame, valid_df: pd.DataFrame) -> Tuple[float, float]:\n    best_th = 0\n    best_cv = 0\n    for th in [0.1 * x for x in range(11)]:\n        all_preds = get_predictions(valid_df, threshold=th)\n\n        cv = 0\n        for i, row in val_targets_df.iterrows():\n            target = row.target\n            preds = all_preds[row.image]\n            val_targets_df.loc[i, th] = map_per_image(target, preds)\n\n        cv = val_targets_df[th].mean()\n\n        print(f\"th={th} cv={cv}\")\n\n        if cv > best_cv:\n            best_th = th\n            best_cv = cv\n\n    print(f\"best_th={best_th}\")\n    print(f\"best_cv={best_cv}\")\n\n    # Adjustment: Since Public lb has nearly 10% 'new_individual' (Be Careful for private LB)\n    val_targets_df[\"is_new_individual\"] = val_targets_df.target == \"new_individual\"\n    val_scores = val_targets_df.groupby(\"is_new_individual\").mean().T\n    val_scores[\"adjusted_cv\"] = val_scores[True] * 0.1 + val_scores[False] * 0.9\n    best_th = val_scores[\"adjusted_cv\"].idxmax()\n    print(f\"best_th_adjusted={best_th}\")\n\n    return best_th, best_cv\n\n\ndef get_predictions(df: pd.DataFrame, threshold: float = 0.2):\n    sample_list = [\"938b7e931166\", \"5bf17305f073\", \"7593d2aee842\", \"7362d7a01d00\", \"956562ff2888\"]\n\n    predictions = {}\n    for i, row in tqdm(df.iterrows(), total=len(df), desc=f\"Creating predictions for threshold={threshold}\"):\n        if row.image in predictions:\n            if len(predictions[row.image]) == 5:\n                continue\n            predictions[row.image].append(row.target)\n        elif row.distances > threshold:\n            predictions[row.image] = [row.target, \"new_individual\"]\n        else:\n            predictions[row.image] = [\"new_individual\", row.target]\n\n    for x in tqdm(predictions):\n        if len(predictions[x]) < 5:\n            remaining = [y for y in sample_list if y not in predictions]\n            predictions[x] = predictions[x] + remaining\n            predictions[x] = predictions[x][:5]\n\n    return predictions\n\n\n# TODO: add types\ndef map_per_image(label, predictions):\n    \"\"\"Computes the precision score of one image.\n\n    Parameters\n    ----------\n    label : string\n            The true label of the image\n    predictions : list\n            A list of predicted elements (order does matter, 5 predictions allowed per image)\n\n    Returns\n    -------\n    score : double\n    \"\"\"\n    try:\n        return 1 / (predictions[:5].index(label) + 1)\n    except ValueError:\n        return 0.0\n\n\ndef create_predictions_df(test_df: pd.DataFrame, best_th: float) -> pd.DataFrame:\n    predictions = get_predictions(test_df, best_th)\n\n    predictions = pd.Series(predictions).reset_index()\n    predictions.columns = [\"image\", \"predictions\"]\n    predictions[\"predictions\"] = predictions[\"predictions\"].apply(lambda x: \" \".join(x))\n\n    return predictions","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:50.200163Z","iopub.execute_input":"2022-04-16T01:09:50.200415Z","iopub.status.idle":"2022-04-16T01:09:50.235298Z","shell.execute_reply.started":"2022-04-16T01:09:50.200375Z","shell.execute_reply":"2022-04-16T01:09:50.234607Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def infer(\n    checkpoint_path: str,\n    checkpoint_path2: str,\n    train_csv_encoded_folded: str = str(TRAIN_CSV_ENCODED_FOLDED_PATH),\n    test_csv: str = str(TEST_CSV_PATH),\n    val_fold: float = 0.0,\n    image_size: int = 256,\n    batch_size: int = 32,\n    num_workers: int = 2,\n    k: int = 50,\n):\n    module = load_eval_module(checkpoint_path, torch.device(\"cuda\"))\n    #module2 = load_eval_module(checkpoint_path2, torch.device(\"cuda\"))\n    train_dl, val_dl, test_dl = load_dataloaders(\n        train_csv_encoded_folded=train_csv_encoded_folded,\n        test_csv=test_csv,\n        val_fold=val_fold,\n        image_size=image_size,\n        batch_size=batch_size,\n        num_workers=num_workers,\n    )\n\n    encoder = load_encoder()\n    \n    #_, train_embeddings2, train_targets2 = get_embeddings(module2, train_dl, encoder, stage=\"train\")\n    #_, val_embeddings2, val_targets2 = get_embeddings(module2, val_dl, encoder, stage=\"val\")\n    #_, test_embeddings2, test_targets2 = get_embeddings(module2, test_dl, encoder, stage=\"test\")\n    \n    #train_image_names, train_embeddings, train_targets = get_embeddings(module,module2, train_dl, encoder, stage=\"train\")\n    #val_image_names, val_embeddings, val_targets = get_embeddings(module,module2, val_dl, encoder, stage=\"val\")\n    #test_image_names, test_embeddings, test_targets = get_embeddings(module,module2, test_dl, encoder, stage=\"test\")\n    train_image_names, train_embeddings, train_targets = get_embeddings(module, train_dl, encoder, stage=\"train\")\n    val_image_names, val_embeddings, val_targets = get_embeddings(module, val_dl, encoder, stage=\"val\")\n    test_image_names, test_embeddings, test_targets = get_embeddings(module, test_dl, encoder, stage=\"test\")\n\n    #ensemble\n\n    #train_embeddings = np.concatenate([train_embeddings, train_embeddings2])\n    #val_embeddings = np.concatenate([val_embeddings, val_embeddings2])\n    #test_embeddings = np.concatenate([test_embeddings, test_embeddings2])\n    \n    #train_targets = np.concatenate([train_targets, train_targets2])\n    #val_targets = np.concatenate([val_targets, val_targets2])\n    #test_targets = np.concatenate([test_targets, test_targets2])\n    #train_embeddings = np.mean(train_embeddings+ train_embeddings2)\n    #val_embeddings = np.mean(val_embeddings+ val_embeddings2)\n    #test_embeddings = np.mean(test_embeddings+ test_embeddings2)\n    \n    D, I = create_and_search_index(module.hparams.embedding_size, train_embeddings, val_embeddings, k)  # noqa: E741\n    print(\"Created index with train_embeddings\")\n\n    val_targets_df = create_val_targets_df(train_targets, val_image_names, val_targets)\n    print(f\"val_targets_df=\\n{val_targets_df.head()}\")\n\n    val_df = create_distances_df(val_image_names, train_targets, D, I, \"val\")\n    print(f\"val_df=\\n{val_df.head()}\")\n\n    best_th, best_cv = get_best_threshold(val_targets_df, val_df)\n    print(f\"val_targets_df=\\n{val_targets_df.describe()}\")\n\n    train_embeddings = np.concatenate([train_embeddings, val_embeddings])\n    train_targets = np.concatenate([train_targets, val_targets])\n    print(\"Updated train_embeddings and train_targets with val data\")\n\n    D, I = create_and_search_index(module.hparams.embedding_size, train_embeddings, test_embeddings, k)  # noqa: E741\n    print(\"Created index with train_embeddings\")\n\n    test_df = create_distances_df(test_image_names, train_targets, D, I, \"test\")\n    print(f\"test_df=\\n{test_df.head()}\")\n\n    predictions = create_predictions_df(test_df, best_th)\n    print(f\"predictions.head()={predictions.head()}\")\n    \n    # Fix missing predictions\n    # From https://www.kaggle.com/code/jpbremer/backfins-arcface-tpu-effnet/notebook\n    \n    public_predictions = pd.read_csv(PUBLIC_SUBMISSION_CSV_PATH)\n    ids_without_backfin = np.load(IDS_WITHOUT_BACKFIN_PATH, allow_pickle=True)\n\n    ids2 = public_predictions[\"image\"][~public_predictions[\"image\"].isin(predictions[\"image\"])]\n\n    predictions = pd.concat(\n        [\n            predictions[~(predictions[\"image\"].isin(ids_without_backfin))],\n            public_predictions[public_predictions[\"image\"].isin(ids_without_backfin)],\n            public_predictions[public_predictions[\"image\"].isin(ids2)],\n        ]\n    )\n    predictions = predictions.drop_duplicates()\n\n    predictions.to_csv(SUBMISSION_CSV_PATH, index=False)\n    ","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:50.238085Z","iopub.execute_input":"2022-04-16T01:09:50.238322Z","iopub.status.idle":"2022-04-16T01:09:50.254927Z","shell.execute_reply.started":"2022-04-16T01:09:50.238297Z","shell.execute_reply":"2022-04-16T01:09:50.254098Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_name = \"convnext_base_384_in22ft1k\"\nimage_size = 384\nbatch_size = 64\nvfold = 4.0\n#train(model_name=model_name,\n      #image_size=image_size,\n     # batch_size=batch_size)\n\ninfer(checkpoint_path=\"../input/cnextfold4/convnext_base_384_in22ft1k_fold4.ckpt\",\n      checkpoint_path2=\"../input/convnext-day323/convnext_base_384_in22ft1k_384.ckpt\",\n      image_size=image_size, batch_size=batch_size,k=100, val_fold = vfold)","metadata":{"execution":{"iopub.status.busy":"2022-04-16T01:09:50.256031Z","iopub.execute_input":"2022-04-16T01:09:50.256404Z","iopub.status.idle":"2022-04-16T01:43:38.340960Z","shell.execute_reply.started":"2022-04-16T01:09:50.256340Z","shell.execute_reply":"2022-04-16T01:43:38.340185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}