Upload src/training/losses.py with huggingface_hub

src/training/losses.py

@@ -0,0 +1,200 @@
+"""
+Custom loss functions for thermal pattern analysis training.
+
+Implements:
+    - ContrastiveLoss  — pushes same-class pairs together, different-class apart
+    - TripletLoss      — anchor / positive / negative margin ranking
+    - CombinedLoss     — weighted sum of both
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ContrastiveLoss(nn.Module):
+    """
+    Contrastive loss (Chopra et al., 2005).
+
+    For a pair of embeddings (e1, e2) with label y ∈ {0, 1}:
+        y=0 → same class   → loss = ½ · D²
+        y=1 → diff class   → loss = ½ · max(0, margin − D)²
+    where D = ‖e1 − e2‖₂.
+    """
+
+    def __init__(self, margin: float = 1.0):
+        super().__init__()
+        self.margin = margin
+
+    def forward(
+        self,
+        embeddings1: torch.Tensor,
+        embeddings2: torch.Tensor,
+        labels: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            embeddings1: (B, D)
+            embeddings2: (B, D)
+            labels:      (B,)  — 0 if same class, 1 if different
+
+        Returns:
+            Scalar loss.
+        """
+        distance = F.pairwise_distance(embeddings1, embeddings2)
+        loss = (
+            (1 - labels) * distance.pow(2)
+            + labels * F.relu(self.margin - distance).pow(2)
+        )
+        return 0.5 * loss.mean()
+
+
+class TripletLoss(nn.Module):
+    """
+    Triplet margin loss with optional hard-negative mining.
+
+    loss = max(0, d(a, p) − d(a, n) + margin)
+    """
+
+    def __init__(self, margin: float = 1.0):
+        super().__init__()
+        self.loss_fn = nn.TripletMarginLoss(margin=margin, p=2)
+
+    def forward(
+        self,
+        anchor: torch.Tensor,
+        positive: torch.Tensor,
+        negative: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            anchor:   (B, D)
+            positive: (B, D)  — same class as anchor
+            negative: (B, D)  — different class from anchor
+
+        Returns:
+            Scalar loss.
+        """
+        return self.loss_fn(anchor, positive, negative)
+
+
+class CombinedLoss(nn.Module):
+    """
+    Weighted combination of Contrastive and Triplet losses,
+    with a standard cross-entropy classification head.
+
+    total = α·contrastive + β·triplet + γ·classification
+    """
+
+    def __init__(
+        self,
+        contrastive_weight: float = 0.3,
+        triplet_weight: float = 0.3,
+        classification_weight: float = 0.4,
+        triplet_margin: float = 1.0,
+        contrastive_margin: float = 1.0,
+    ):
+        super().__init__()
+        self.contrastive_weight = contrastive_weight
+        self.triplet_weight = triplet_weight
+        self.classification_weight = classification_weight
+
+        self.contrastive_loss = ContrastiveLoss(margin=contrastive_margin)
+        self.triplet_loss = TripletLoss(margin=triplet_margin)
+        self.classification_loss = nn.CrossEntropyLoss()
+
+    @classmethod
+    def from_config(cls, config) -> "CombinedLoss":
+        """Construct from a Config object."""
+        loss_cfg = config.training.loss
+        return cls(
+            contrastive_weight=loss_cfg.contrastive_weight,
+            triplet_weight=loss_cfg.triplet_weight,
+            classification_weight=1.0 - loss_cfg.contrastive_weight - loss_cfg.triplet_weight,
+            triplet_margin=loss_cfg.triplet_margin,
+        )
+
+    def forward(
+        self,
+        embeddings: torch.Tensor,
+        labels: torch.Tensor,
+        logits: torch.Tensor | None = None,
+    ) -> dict:
+        """
+        Compute the combined loss.
+
+        Uses in-batch pair and triplet mining for efficiency.
+
+        Args:
+            embeddings: (B, D)
+            labels:     (B,) integer class labels
+            logits:     (B, num_classes) or None
+
+        Returns:
+            dict with total_loss, contrastive, triplet, classification.
+        """
+        total = torch.tensor(0.0, device=embeddings.device)
+        result = {}
+
+        # ------- Contrastive: generate in-batch pairs -------
+        B = embeddings.size(0)
+        if B >= 2:
+            idx = torch.randperm(B, device=embeddings.device)
+            e1, e2 = embeddings, embeddings[idx]
+            pair_labels = (labels != labels[idx]).float()
+
+            c_loss = self.contrastive_loss(e1, e2, pair_labels)
+            total = total + self.contrastive_weight * c_loss
+            result["contrastive"] = c_loss.item()
+
+        # ------- Triplet: mine anchor / pos / neg -------
+        anchors, positives, negatives = self._mine_triplets(embeddings, labels)
+        if anchors is not None:
+            t_loss = self.triplet_loss(anchors, positives, negatives)
+            total = total + self.triplet_weight * t_loss
+            result["triplet"] = t_loss.item()
+
+        # ------- Classification -------
+        if logits is not None:
+            cls_loss = self.classification_loss(logits, labels)
+            total = total + self.classification_weight * cls_loss
+            result["classification"] = cls_loss.item()
+
+        result["total_loss"] = total
+        return result
+
+    @staticmethod
+    def _mine_triplets(
+        embeddings: torch.Tensor, labels: torch.Tensor
+    ) -> tuple:
+        """Simple in-batch triplet mining."""
+        unique_labels = labels.unique()
+        if len(unique_labels) < 2:
+            return None, None, None
+
+        anchors, positives, negatives = [], [], []
+
+        for label in unique_labels:
+            mask_pos = labels == label
+            mask_neg = labels != label
+
+            pos_idx = mask_pos.nonzero(as_tuple=True)[0]
+            neg_idx = mask_neg.nonzero(as_tuple=True)[0]
+
+            if len(pos_idx) < 2 or len(neg_idx) < 1:
+                continue
+
+            for i in range(min(len(pos_idx) - 1, 4)):  # limit per class
+                anchors.append(embeddings[pos_idx[i]])
+                positives.append(embeddings[pos_idx[i + 1]])
+                neg_i = neg_idx[torch.randint(len(neg_idx), (1,)).item()]
+                negatives.append(embeddings[neg_i])
+
+        if not anchors:
+            return None, None, None
+
+        return (
+            torch.stack(anchors),
+            torch.stack(positives),
+            torch.stack(negatives),
+        )