{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":91451,"databundleVersionId":11223220,"sourceType":"competition"},{"sourceId":238980352,"sourceType":"kernelVersion"}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, random\nfrom collections import defaultdict\nfrom typing import Any\nfrom tqdm import tqdm\nfrom sklearn.preprocessing import LabelEncoder\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import train_test_split\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader, RandomSampler\nimport functools\nimport operator\nimport torch.nn.functional as F\n\n\ndef split_reid_data(\n    df: pd.DataFrame,\n    train_ratio: float = 0.6,\n    val_ratio: float = 0.5,\n    known_ratio: float = 0.8,\n    group_col: str = \"identity\",\n    seed: int = 42,\n) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:\n    \"\"\"\n    Splits the dataset into training, validation, and test sets for open-set re-identification tasks.\n    Simplified method to get a training/validation split\n    \"\"\"\n    # 1. Split into Train and (Temp) Test sets based on individuals (groups)\n    unique_ids = df[group_col].unique()\n    train_ids, temp_test_ids = train_test_split(\n        unique_ids, train_size=train_ratio, random_state=seed\n    )\n\n    # 2. Split (Temp) Test into Validation and Final Test sets (unknown individuals).\n    val_unknown_ids, test_unknown_ids = train_test_split(\n        temp_test_ids, train_size=val_ratio, random_state=seed + 1\n    )\n\n    return (\n        df[df[group_col].isin(train_ids)], \n        df[df[group_col].isin(val_unknown_ids)], \n        df[df[group_col].isin(test_unknown_ids)]\n    )\n\n\nclass EmbeddingDataset(Dataset):\n    def __init__(self, metadata: pd.DataFrame,):\n        self.metadata = self._filter_metadata(metadata)\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        display(self.metadata.groupby(\"dataset\").count())\n        self.label_encoder = LabelEncoder()\n        self.label_encoder.fit(self.metadata.identity)\n    \n    def _filter_metadata(self, metadata):\n        # only do triplet learning on the database images that have at least 2 images\n        metadata = metadata.loc[metadata['split'] == 'database']\n        dfs = []\n        for identity in metadata.identity.unique():\n            id_metadata = metadata.loc[metadata.identity == identity].copy()\n            if id_metadata.shape[0] >= 2:\n                dfs.append(id_metadata)\n        return pd.concat(dfs, axis=0).reset_index()\n        \n    def __len__(self) -> int:\n        return len(self.metadata)\n\n    def __getitem__(self, idx: int) -> tuple:\n        row = self.metadata.iloc[idx]\n        embedding = torch.from_numpy(row.embeddings).to(self.device)\n        idx = torch.tensor(idx).to(device)\n        label = torch.tensor(self.label_encoder.transform([row.identity])[0]).to(device)\n        return embedding, idx, label\n\n    @functools.cache\n    def _filter_positive(self, idx):\n        df = self.metadata\n        anchor = df.iloc[idx]\n        # find all the rows that match the condition and sample 1\n        cond = [\n            (df.identity == anchor.identity),\n            (df.image_id != anchor.image_id)\n        ]\n        # we only need the index from these so we try to reduce the cache\n        return df[functools.reduce(operator.__and__, cond)].index\n\n    @functools.cache\n    def _filter_negative(self, idx, same_species=False):\n        \"\"\"find a random row not in the current identity\"\"\"\n        df = self.metadata\n        anchor = df.iloc[idx]\n        cond = [(df.identity != anchor.identity)]\n        if same_species:\n            cond.append((df.dataset == anchor.dataset))\n        # we only need the index from these so we try to reduce the cache\n        return df[functools.reduce(operator.__and__, cond)].index\n    \n    def sample_positive(self, idx, n=1):\n        \"\"\"find a row with the same label but different index\"\"\"\n        return random.choices(self._filter_positive(idx), k=n)\n\n    def sample_negative(self, idx, n=1, same_species=False):\n        \"\"\"find a random row not in the current identity\"\"\"\n        return random.choices(self._filter_negative(idx), k=n)\n\n\nclass ProjectionHead(nn.Module):\n    \"\"\"\n    Simple projection head to transform embeddings.\n    \n    Tried out nonlinearity (GELU) but not sure if this is actually\n    appropriate given that dino features give a good representational basis.\n    Normalization is useful because it links the cosine distance to \n    euclidean distance via inner product.\n    \"\"\"\n\n    def __init__(self, input_dim, output_dim, hidden_dim=None):\n        super().__init__()\n        if hidden_dim is None:\n            hidden_dim = input_dim\n    \n        self.projection = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.GELU(),\n            nn.Linear(hidden_dim, output_dim)\n        )\n\n    def forward(self, x):\n        z = self.projection(x)\n        z = F.normalize(z, p=2, dim=-1) \n        return z\n\n\nclass TripletLoss(nn.Module):\n    \"\"\"\n    Triplet loss with a margin.\n    Takes embeddings of an anchor, a positive and a negative sample.\n    \"\"\"\n\n    def __init__(self, margin=1.0):\n        super().__init__()\n        self.margin = margin\n\n    def forward(self, anchor, positive, negative):\n        pos_dist = torch.pairwise_distance(anchor, positive)\n        neg_dist = torch.pairwise_distance(anchor, negative)\n        loss = torch.relu(pos_dist - neg_dist + self.margin)\n        return loss.mean()\n\n\nclass MRLTripletLoss(nn.Module):\n    \"\"\"\n    Matryoshka Representation Learning with Triplet Loss.\n    Applies triplet loss at multiple embedding dimensionalities (prefixes).\n    \"\"\"\n    def __init__(self, \n                 margin: float = 1.0, \n                 nested_dims: list = [],\n                 loss_weights: list = None,\n                ):\n        super().__init__()\n        self.margin = margin\n        self.base_triplet_loss = TripletLoss(margin=margin)\n        self.nested_dims = nested_dims\n        \n        # Weights for each loss component (full + nested).\n        # If not provided, defaults to equal weighting (1.0 for each).\n        num_losses = 1 + len(self.nested_dims)\n        self.loss_weights = loss_weights\n        if self.loss_weights is None:\n            self.loss_weights = [1.0] * num_losses\n        assert len(self.loss_weights) == num_losses\n        self.loss_weight_sum = sum(self.loss_weights)\n\n    def forward(self, anchor_full, positive_full, negative_full):\n        total_loss = 0.0\n        \n        loss_full = self.base_triplet_loss(anchor_full, positive_full, negative_full)\n        total_loss += self.loss_weights[0] * loss_full\n        \n        for i, d_prefix in enumerate(self.nested_dims):\n            if d_prefix >= anchor_full.shape[-1]:\n                continue\n\n            anchor_prefix = anchor_full[:, :d_prefix]\n            positive_prefix = positive_full[:, :d_prefix]\n            negative_prefix = negative_full[:, :d_prefix]\n            \n            loss_prefix = self.base_triplet_loss(anchor_prefix, positive_prefix, negative_prefix)\n            total_loss += self.loss_weights[i+1] * loss_prefix\n            \n        return total_loss / self.loss_weight_sum\n\n\ndef train_batch_semi_hard_negative(\n    metadata,\n    projection_head,\n    batch_size=200,\n    epochs=100,\n    margin=1.0,\n    nested_dims=[],\n    loss_weights=None,\n    learning_rate=5e-4,\n    warmup_epochs=10,\n    triplets_per_anchor=1,\n):\n    train_df, val_df, _ = split_reid_data(metadata)\n\n    metadatas = {\n        \"train\": train_df,\n        \"val\": val_df,\n    }\n    datasets = {}\n    dataloaders = {}\n    for step in [\"train\", \"val\"]:\n        datasets[step] = EmbeddingDataset(metadatas[step])\n        dataloaders[step] = DataLoader(\n            datasets[step],\n            batch_size=batch_size,\n            sampler=RandomSampler(datasets[step]),\n        )\n\n    optimizer = optim.Adam(projection_head.parameters(), lr=learning_rate)\n    warmup_scheduler = optim.lr_scheduler.LinearLR(\n        optimizer,\n        start_factor=0.1,\n        end_factor=1.0,\n        total_iters=warmup_epochs\n    )\n    cosine_scheduler = optim.lr_scheduler.CosineAnnealingLR(\n        optimizer,\n        T_max=epochs - warmup_epochs,\n        eta_min=1e-7\n    )\n    scheduler = optim.lr_scheduler.SequentialLR(\n        optimizer,\n        schedulers=[warmup_scheduler, cosine_scheduler],\n        milestones=[warmup_epochs]\n    )\n\n    device = next(projection_head.parameters()).device\n    criterion = MRLTripletLoss(margin=margin, nested_dims=nested_dims, loss_weights=loss_weights)\n\n    for epoch in tqdm(range(epochs)):\n        for step in [\"train\", \"val\"]:\n            if step == \"train\":\n                projection_head.train()\n            else:\n                projection_head.eval()\n            dataset = datasets[step]\n            dataloader = dataloaders[step]\n\n            running_loss = 0.0\n            batch_count = 0\n            epoch_triplets = 0\n    \n            for batch_embeddings, batch_indices, batch_labels in dataloader:\n                @functools.cache\n                def _embed(idx):\n                    emb, _, _ = dataset[idx]\n                    return projection_head(emb).unsqueeze(0)\n    \n                batch_size = batch_embeddings.size(0)\n                if batch_size <= 1:\n                    continue\n    \n                batch_embeddings = projection_head(batch_embeddings)\n                batch_pairwise_dist = torch.cdist(\n                    batch_embeddings, batch_embeddings, p=2\n                )\n    \n                # Process each anchor in the batch\n                batch_loss = 0\n                batch_triplets = 0\n    \n                for anchor_embedding, anchor_index, anchor_label, anchor_dist in zip(\n                    batch_embeddings, batch_indices, batch_labels, batch_pairwise_dist\n                ):\n                    anchor_embedding = anchor_embedding.unsqueeze(0)\n                    anchor_index = int(anchor_index)\n                    \n                    # we're going to generate triplets in the following way\n                    # first we get a random positive of indices\n                    # then we do batch semi-hard mining by looking for everything within the margin\n                    # skip if it doesn't exist\n                    for _ in range(triplets_per_anchor):\n                        pos_idx = dataset.sample_positive(anchor_index)[0]\n                        pos_embedding = _embed(pos_idx)\n                        # distance positive -> dp\n                        dp = torch.pairwise_distance(anchor_embedding, pos_embedding)\n    \n        \n                        # BATCH SEMI-HARD NEGATIVE MINING\n                        # Distances from the current anchor (projected) to all other *projected* samples in THIS BATCH\n                        # Mask for negatives within the current batch (different true label than anchor)                    \n                        \n                        # Semi-hard condition mask\n                        semi_hard_mask = (\n                            # labels that dont match the anchor\n                            (batch_labels != anchor_label)\n                            # dp < dn < dp+margin\n                            & (anchor_dist > dp)\n                            & (anchor_dist < (dp + criterion.margin))\n                        )\n    \n                        # skip because no semi-hard\n                        if not torch.any(semi_hard_mask):\n                            continue \n        \n                        # Get indices (within the current batch) of these semi-hard negatives\n                        semi_hard_indices = torch.where(semi_hard_mask)[0]\n                        \n                        # Select one semi-hard negative (e.g., the hardest of the semi-hard, or random)\n                        # To get the \"hardest\" of the semi-hard (closest to anchor among semi-hard):\n                        smallest_neg = torch.min(anchor_dist[semi_hard_indices])\n                        # Find which batch index corresponds to this distance among the semi-hard ones\n                        neg_indices = torch.where(semi_hard_mask & (anchor_dist == smallest_neg))[0]\n                        \n                        if neg_indices.numel() == 0: continue\n                        neg_idx = neg_indices[0]\n        \n                        neg_embedding = batch_embeddings[neg_idx].unsqueeze(0)\n                        triplet_loss = criterion(\n                            anchor_embedding, pos_embedding, neg_embedding\n                        )\n                        batch_loss += triplet_loss\n                        batch_triplets += 1\n                        epoch_triplets += 1\n    \n                # Skip batch if no valid triplets\n                if batch_triplets == 0:\n                    continue\n    \n                # Normalize loss by number of triplets\n                batch_loss = batch_loss / batch_triplets\n                if step == \"train\":\n                    optimizer.zero_grad()\n                    batch_loss.backward()\n                    optimizer.step()\n    \n                # Update statistics\n                running_loss += batch_loss.item()\n                batch_count += 1\n    \n            # Print epoch statistics\n            if batch_count > 0:\n                current_lr = optimizer.param_groups[0]['lr']\n                epoch_loss = running_loss / batch_count\n                print(f\"Step {step}, Epoch {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}, Triplet Count: {epoch_triplets}, LR: {current_lr}\")\n            else:\n                print(f\"Step {step}, Epoch {epoch+1}/{epochs}, No valid triplets found\")\n\n\n        # done with loop so update scheduler\n        scheduler.step()\n        \n        # save intermediate heads\n        if epoch > 0 and epoch % 10 == 0:\n            name =  f\"head_epoch={epoch:03d}.pt\"\n            print(f\"saving {name}\")\n            torch.save(projection_head.state_dict(), name)\n            \n    return projection_head\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T05:58:41.651240Z","iopub.execute_input":"2025-05-13T05:58:41.651533Z","iopub.status.idle":"2025-05-13T05:58:41.680796Z","shell.execute_reply.started":"2025-05-13T05:58:41.651511Z","shell.execute_reply":"2025-05-13T05:58:41.680082Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"metadata = pd.read_csv('/kaggle/input/animal-clef-2025/metadata.csv')\nembeddings_df = pd.read_parquet(\"/kaggle/input/preprocess-triplet-embedding/embeddings.parquet\")\nmerged_df = pd.merge(metadata, embeddings_df, on='image_id', how='inner')\ndisplay(merged_df.head(3))\n\ndim = 256\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nproj_head = ProjectionHead(768, dim).to(device)\nhead = train_batch_semi_hard_negative(\n    merged_df, \n    projection_head=proj_head,\n    epochs=100,\n    nested_dims=[128, 64, 32, 16, 2],\n    loss_weights=[1, 128/dim, 64/dim, 32/dim, 16/dim, 2/dim],\n    # disable gamma\n    warmup_epochs=10,\n    learning_rate=1e-3,\n)\n\n# write the output\ntorch.save(head.state_dict(), '/kaggle/working/head.pt')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T05:58:41.940115Z","iopub.execute_input":"2025-05-13T05:58:41.940385Z","iopub.status.idle":"2025-05-13T06:05:55.993197Z","shell.execute_reply.started":"2025-05-13T05:58:41.940366Z","shell.execute_reply":"2025-05-13T06:05:55.992002Z"}},"outputs":[],"execution_count":null}]}