{"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 ⚡ BackFin ConvNeXt ArcFace\n\nLet's train [`timm`](https://github.com/rwightman/pytorch-image-models) models with [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- [backfintfrecrods](https://www.kaggle.com/datasets/jpbremer/backfintfrecords) dataset\n\n## Nice Lightning `Trainer` Flags to try\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.701 (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":"!pip install -q timm faiss-gpu\n!pip install -q happywhale -U -f ../input/-pytorch-lightning-happywhale-pkg\n!pip install -q Pillow==9.0.0\n!pip uninstall -y torchtext\n!pip list | grep torch\n!pip list | grep lightning\n!pip list | grep happywhale","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-04-18T13:21:57.761467Z","iopub.execute_input":"2022-04-18T13:21:57.762081Z","iopub.status.idle":"2022-04-18T13:22:35.399187Z","shell.execute_reply.started":"2022-04-18T13:21:57.762043Z","shell.execute_reply":"2022-04-18T13:22:35.398351Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import math\nfrom typing import Callable, Dict, Optional, 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, Dataset\nfrom tqdm.notebook import tqdm\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.preprocessing import normalize, LabelEncoder","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-04-18T13:22:35.401514Z","iopub.execute_input":"2022-04-18T13:22:35.401798Z","iopub.status.idle":"2022-04-18T13:22:35.410902Z","shell.execute_reply.started":"2022-04-18T13:22:35.401758Z","shell.execute_reply":"2022-04-18T13:22:35.408341Z"},"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\"\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 = OUTPUT_DIR / \"checkpoints\"\nSUBMISSION_CSV_PATH = OUTPUT_DIR / \"submission.csv\"","metadata":{"execution":{"iopub.status.busy":"2022-04-18T13:22:35.412188Z","iopub.execute_input":"2022-04-18T13:22:35.412710Z","iopub.status.idle":"2022-04-18T13:22:35.422624Z","shell.execute_reply.started":"2022-04-18T13:22:35.412673Z","shell.execute_reply":"2022-04-18T13:22:35.421842Z"},"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 str(dir / id)","metadata":{"execution":{"iopub.status.busy":"2022-04-18T13:22:35.425986Z","iopub.execute_input":"2022-04-18T13:22:35.427482Z","iopub.status.idle":"2022-04-18T13:22:35.434204Z","shell.execute_reply.started":"2022-04-18T13:22:35.427444Z","shell.execute_reply":"2022-04-18T13:22:35.433323Z"},"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 = train_df.astype({'kfold': 'int8'})\ntrain_df.to_csv(TRAIN_CSV_ENCODED_FOLDED_PATH, index=False)\ndisplay(train_df.head())","metadata":{"execution":{"iopub.status.busy":"2022-04-18T13:22:35.435538Z","iopub.execute_input":"2022-04-18T13:22:35.436029Z","iopub.status.idle":"2022-04-18T13:22:36.449063Z","shell.execute_reply.started":"2022-04-18T13:22:35.435983Z","shell.execute_reply":"2022-04-18T13:22:36.448382Z"},"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\ntest_df.to_csv(TEST_CSV_PATH, index=False)\ntest_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-04-18T13:22:36.450330Z","iopub.execute_input":"2022-04-18T13:22:36.451066Z","iopub.status.idle":"2022-04-18T13:22:36.835147Z","shell.execute_reply.started":"2022-04-18T13:22:36.451026Z","shell.execute_reply":"2022-04-18T13:22:36.834462Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Lightning DataModule","metadata":{}},{"cell_type":"code","source":"from happywhale.utils.dataset import HappyWhaleDataset\n\nclass LitDataModule(pl.LightningDataModule):\n    def __init__(\n        self,\n        train_csv_encoded_folded: str,\n        test_csv: str,\n        image_size: int,\n        batch_size: int,\n        num_workers: int,\n        val_split: float = 0.1,\n        val_fold: float = None,\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        self.num_classes = len(set(\n            list(self.train_df[\"individual_id\"].values) + list(self.test_df[\"individual_id\"].values)\n        ))\n        \n    def setup(self, stage: Optional[str] = None):\n        if stage == \"fit\" or stage is None:\n            nb_train = int(self.hparams.val_split * len(self.train_df))\n            # Split train df using fold\n            if self.hparams.val_fold is not None:\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            else:\n                train_df = self.train_df[:-nb_train].reset_index(drop=True)\n                val_df = self.train_df[-nb_train:].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-18T13:22:36.837480Z","iopub.execute_input":"2022-04-18T13:22:36.837885Z","iopub.status.idle":"2022-04-18T13:22:36.852224Z","shell.execute_reply.started":"2022-04-18T13:22:36.837846Z","shell.execute_reply":"2022-04-18T13:22:36.851498Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Lightning Module","metadata":{}},{"cell_type":"code","source":"from happywhale.utils.arc_margin_product import ArcMarginProduct\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        batch_size: int,\n        epochs: int\n    ):\n        super().__init__()\n\n        self.save_hyperparameters()\n\n        self.model = timm.create_model(model_name, pretrained=pretrained, drop_rate=drop_rate)\n        self.embedding = nn.Linear(self.model.get_classifier().in_features, embedding_size)\n        self.model.reset_classifier(num_classes=0, global_pool=\"avg\")\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\n        self.loss_fn = F.cross_entropy\n\n    def forward(self, images: torch.Tensor) -> torch.Tensor:\n        features = self.model(images)\n        embeddings = self.embedding(features)\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        return loss","metadata":{"execution":{"iopub.status.busy":"2022-04-18T13:22:36.853815Z","iopub.execute_input":"2022-04-18T13:22:36.854594Z","iopub.status.idle":"2022-04-18T13:22:36.871795Z","shell.execute_reply.started":"2022-04-18T13:22:36.854548Z","shell.execute_reply":"2022-04-18T13:22:36.870986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"from pytorch_lightning.loggers import CSVLogger\nfrom pytorch_lightning.callbacks import StochasticWeightAveraging\n\n\ndef train(\n    train_csv_encoded_folded: str = str(TRAIN_CSV_ENCODED_FOLDED_PATH),\n    test_csv: str = str(TEST_CSV_PATH),\n    # val_fold: int = 0,\n    val_split: float = 0.05,\n    image_size: int = 256,\n    batch_size: int = 64,\n    num_workers: int = 4,\n    model_name: str = \"tf_efficientnet_b0\",\n    pretrained: bool = True,\n    drop_rate: float = 0.0,\n    embedding_size: int = 512,\n    arc_s: float = 30.0,\n    arc_m: float = 0.5,\n    arc_easy_margin: bool = False,\n    arc_ls_eps: float = 0.0,\n    optimizer: str = \"adamw\",\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 = True,\n    # fast_dev_run: bool = False,\n    gpus: int = 1,\n    max_epochs: int = 10,\n    precision: int = 16,\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_split=val_split,\n        image_size=image_size,\n        batch_size=batch_size,\n        num_workers=num_workers,\n    )\n    \n    datamodule.setup()\n    print(datamodule.num_classes)\n    len_train_dl = len(datamodule.train_dataloader())\n\n    model = LitModule(\n        model_name=model_name,\n        pretrained=pretrained,\n        drop_rate=drop_rate,\n        embedding_size=embedding_size,\n        num_classes=datamodule.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        batch_size=batch_size,\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    swa = StochasticWeightAveraging(swa_epoch_start=0.6)\n    logger = CSVLogger(save_dir='logs/')\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        logger=logger,\n        callbacks=[model_checkpoint],\n        # deterministic=True,\n        # fast_dev_run=fast_dev_run,\n        gpus=gpus,\n        max_epochs=max_epochs,\n        precision=precision,\n        # limit_train_batches=0.1,\n        # limit_val_batches=0.1,\n    )\n\n    # trainer.tune(module, datamodule=datamodule)\n\n    trainer.fit(model, datamodule=datamodule)\n    \n    return model, trainer","metadata":{"execution":{"iopub.status.busy":"2022-04-18T13:22:36.873261Z","iopub.execute_input":"2022-04-18T13:22:36.873673Z","iopub.status.idle":"2022-04-18T13:22:36.889431Z","shell.execute_reply.started":"2022-04-18T13:22:36.873636Z","shell.execute_reply":"2022-04-18T13:22:36.888793Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"https://github.com/rwightman/pytorch-image-models/blob/master/results/results-imagenet.csv","metadata":{}},{"cell_type":"code","source":"MODEL_NAME = \"convnext_tiny\"\nIMAGE_SIZE = 224\nBATCH_SIZE = 48\n\nmodel, trainer = train(model_name=MODEL_NAME, image_size=IMAGE_SIZE, batch_size=BATCH_SIZE)","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-04-18T13:22:36.891864Z","iopub.execute_input":"2022-04-18T13:22:36.892126Z","iopub.status.idle":"2022-04-18T13:31:12.568656Z","shell.execute_reply.started":"2022-04-18T13:22:36.892091Z","shell.execute_reply":"2022-04-18T13:31:12.567825Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport seaborn as sn\n\nmetrics = pd.read_csv(f'{trainer.logger.log_dir}/metrics.csv')\ndel metrics[\"step\"]\nmetrics.set_index(\"epoch\", inplace=True)\ndisplay(metrics.dropna(axis=1, how=\"all\").head())\ng = sn.relplot(data=metrics, kind=\"line\")\nplt.gcf().set_size_inches(12, 4)\nplt.grid()","metadata":{"execution":{"iopub.status.busy":"2022-04-18T13:31:12.570475Z","iopub.execute_input":"2022-04-18T13:31:12.570944Z","iopub.status.idle":"2022-04-18T13:31:13.265695Z","shell.execute_reply.started":"2022-04-18T13:31:12.570897Z","shell.execute_reply":"2022-04-18T13:31:13.264994Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"from happywhale.inference.infer import (\n    load_eval_module,\n    load_encoder,\n    # get_embeddings,\n    create_and_search_index,\n    create_val_targets_df,\n    create_distances_df,\n    get_best_threshold,\n    create_predictions_df,\n)\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    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\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\n        all_image_names.append(image_names)\n        all_embeddings.append(embeddings.cpu().numpy())\n        all_targets.append(targets.cpu().numpy())\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","metadata":{"execution":{"iopub.status.busy":"2022-04-18T13:31:13.266897Z","iopub.execute_input":"2022-04-18T13:31:13.268444Z","iopub.status.idle":"2022-04-18T13:31:14.583324Z","shell.execute_reply.started":"2022-04-18T13:31:13.268400Z","shell.execute_reply":"2022-04-18T13:31:14.582500Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from happywhale.settings import IDS_WITHOUT_BACKFIN_PATH, PUBLIC_SUBMISSION_CSV_PATH\n\ndef infer(\n    checkpoint_path: 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 = 64,\n    num_workers: int = 2,\n    k: int = 50,\n):\n    module = load_eval_module(checkpoint_path, torch.device(\"cuda\"), lit_module_cls=LitModule)\n\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(ENCODER_CLASSES_PATH)\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    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, adjust_th=True)\n    print(\"val_targets_df:\")\n    display(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    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    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[[\"image\",\"predictions\"]].drop_duplicates()\n    predictions.to_csv(SUBMISSION_CSV_PATH, index=False)","metadata":{"execution":{"iopub.status.busy":"2022-04-18T15:03:25.890240Z","iopub.execute_input":"2022-04-18T15:03:25.890625Z","iopub.status.idle":"2022-04-18T15:03:25.906097Z","shell.execute_reply.started":"2022-04-18T15:03:25.890584Z","shell.execute_reply":"2022-04-18T15:03:25.905277Z"},"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"infer(checkpoint_path=CHECKPOINTS_DIR / f\"{MODEL_NAME}_{IMAGE_SIZE}.ckpt\", image_size=IMAGE_SIZE, batch_size=BATCH_SIZE)","metadata":{"execution":{"iopub.status.busy":"2022-04-18T15:03:29.178001Z","iopub.execute_input":"2022-04-18T15:03:29.178310Z","iopub.status.idle":"2022-04-18T15:19:57.024683Z","shell.execute_reply.started":"2022-04-18T15:03:29.178260Z","shell.execute_reply":"2022-04-18T15:19:57.023751Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!head submission.csv","metadata":{"execution":{"iopub.status.busy":"2022-04-18T15:19:57.026976Z","iopub.execute_input":"2022-04-18T15:19:57.027332Z","iopub.status.idle":"2022-04-18T15:19:57.862293Z","shell.execute_reply.started":"2022-04-18T15:19:57.027263Z","shell.execute_reply":"2022-04-18T15:19:57.861356Z"},"trusted":true},"execution_count":null,"outputs":[]}]}