{"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\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\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 (relu) 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):\n        super().__init__()\n        self.fc = nn.Linear(input_dim, output_dim)\n\n    def forward(self, x):\n        z = self.fc(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_with_online_triplet_generation(\n    metadata,\n    projection_head,\n    batch_size=200,\n    epochs=50,\n    margin=1.0,\n    learning_rate=5e-4,\n    lr_step_size=25,\n    lr_gamma=0.1,\n    triplets_per_anchor=1,\n):\n    \"\"\"\n    Train with online triplet generation using a standard DataLoader.\n\n    Args:\n        dataset: Instance of EmbeddingDataset\n        projection_head: projection head to train on top of embeddings\n        batch_size: Batch size for DataLoader\n        epochs: Number of training epochs\n        margin: Margin for triplet loss\n        learning_rate: Learning rate for optimizer\n        triplets_per_anchor: Number of triplets to generate per anchor\n    \"\"\"\n    # Create standard DataLoader with RandomSampler\n    dataset = EmbeddingDataset(metadata)\n    dataloader = DataLoader(\n        dataset,\n        batch_size=batch_size,\n        sampler=RandomSampler(dataset),\n    )\n\n    optimizer = optim.Adam(projection_head.parameters(), lr=learning_rate)\n    # scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=lr_factor, patience=lr_patience, verbose=True, min_lr=1e-7) \n    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=lr_step_size, gamma=lr_gamma) \n    device = next(projection_head.parameters()).device\n\n    # Initialize triplet loss\n    criterion = TripletLoss(margin=margin)\n\n    # Training loop\n    for epoch in tqdm(range(epochs)):\n        running_loss = 0.0\n        batch_count = 0\n\n        # Process batches\n        for batch_embeddings, batch_indices, _ in dataloader:\n            batch_size = batch_embeddings.size(0)\n\n            # Skip small batches\n            if batch_size <= 1:\n                continue\n\n            batch_embeddings = projection_head(batch_embeddings)\n            # NOTE: we want this function to be reset every batch, since\n            # we want to be embedding with the updated head online\n            @functools.cache\n            def _embed(idx):\n                emb, _, _ = dataset[idx]\n                return projection_head(emb).unsqueeze(0)\n\n            # Process each anchor in the batch\n            batch_loss = 0\n            batch_triplets = 0\n\n            for anchor_embedding, anchor_index in zip(batch_embeddings, batch_indices):\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 get a random negative set of indices\n                # then we also get random negative set of indices with the same species\n                sample_size = triplets_per_anchor\n                positive_indices = dataset.sample_positive(anchor_index, sample_size)\n                negative_all_indices = dataset.sample_negative(anchor_index, sample_size)\n\n                # the (pos, neg) pairs of indices from our sampled stuff\n                pairs = [*zip(positive_indices, negative_all_indices)]\n\n                for pos_idx, neg_idx in pairs:\n                    # Compute loss for this triplet\n                    pos_embedding = _embed(pos_idx)\n                    neg_embedding = _embed(neg_idx)\n                    triplet_loss = criterion(\n                        anchor_embedding, pos_embedding, neg_embedding\n                    )\n                    batch_loss += triplet_loss\n                    batch_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            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        # update the scheduler as needed\n        scheduler.step()\n        \n        # Print epoch statistics\n        if batch_count > 0:\n            epoch_loss = running_loss / batch_count\n            print(f\"Epoch {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}\")\n        else:\n            print(f\"Epoch {epoch+1}/{epochs}, No valid triplets found\")\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\n\ndef train_batch_semi_hard_negative(\n    metadata,\n    projection_head,\n    batch_size=200,\n    epochs=50,\n    margin=1.0,\n    nested_dims=[64, 32, 16, 2],\n    loss_weights=[1, 64/128, 32/128, 16/128, 2/128],\n    learning_rate=5e-4,\n    lr_step_size=25,\n    lr_gamma=0.1,\n    triplets_per_anchor=1,\n):\n    dataset = EmbeddingDataset(metadata)\n    dataloader = DataLoader(\n        dataset,\n        batch_size=batch_size,\n        sampler=RandomSampler(dataset),\n    )\n\n    optimizer = optim.Adam(projection_head.parameters(), lr=learning_rate)\n    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=lr_step_size, gamma=lr_gamma) \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        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            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        # update the scheduler as needed\n        scheduler.step()\n        \n        # Print epoch statistics\n        if batch_count > 0:\n            epoch_loss = running_loss / batch_count\n            print(f\"Epoch {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}, Triplet Count: {epoch_triplets}\")\n        else:\n            print(f\"Epoch {epoch+1}/{epochs}, No valid triplets found\")\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-11T20:30:33.567217Z","iopub.execute_input":"2025-05-11T20:30:33.567726Z","iopub.status.idle":"2025-05-11T20:30:33.593687Z","shell.execute_reply.started":"2025-05-11T20:30:33.567698Z","shell.execute_reply":"2025-05-11T20:30:33.592842Z"}},"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\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nproj_head = ProjectionHead(768, 128).to(device)\nhead = train_batch_semi_hard_negative(\n    merged_df, \n    projection_head=proj_head,\n    epochs=100,\n)\n\n# write the output\ntorch.save(head.state_dict(), '/kaggle/working/head.pt')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-11T20:38:25.374987Z","iopub.execute_input":"2025-05-11T20:38:25.375654Z","iopub.status.idle":"2025-05-11T20:44:08.319401Z","shell.execute_reply.started":"2025-05-11T20:38:25.375628Z","shell.execute_reply":"2025-05-11T20:44:08.318795Z"}},"outputs":[],"execution_count":null}]}