first commit

models/losses.py

@@ -0,0 +1,78 @@
+"""Loss terms: Hungarian-matched classification (cross-entropy) + curve regression
+(Smooth L1) + endpoint regression. Reasoning-module auxiliary loss is Phase 2.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.head import eval_curve
+from models.matcher import hungarian_match
+
+
+def compute_loss(
+    outputs: dict,
+    targets: dict,
+    sample_ys: torch.Tensor,
+    cls_weight: float = 2.0,
+    reg_weight: float = 5.0,
+    endpoint_weight: float = 1.0,
+    bg_class_weight: float = 0.2,
+) -> dict:
+    cls_logits = outputs["cls_logits"]        # (B, N, 2)
+    curve_coeffs = outputs["curve_coeffs"]     # (B, N, 4)
+    endpoints = outputs["endpoints"]           # (B, N, 2)
+
+    target_xs = targets["target_xs"]
+    target_valid_mask = targets["target_valid_mask"]
+    target_lane_valid = targets["target_lane_valid"]
+    target_endpoints = targets["target_endpoints"]
+
+    B, N, _ = cls_logits.shape
+    device = cls_logits.device
+
+    matches = hungarian_match(
+        cls_logits, curve_coeffs, target_xs, target_valid_mask, target_lane_valid,
+        sample_ys, cls_weight=cls_weight, reg_weight=reg_weight,
+    )
+
+    cls_targets = torch.zeros(B, N, dtype=torch.long, device=device)
+    reg_losses = []
+    endpoint_losses = []
+
+    for b, (pred_idx, gt_idx) in enumerate(matches):
+        if pred_idx.numel() == 0:
+            continue
+        pred_idx = pred_idx.to(device)
+        gt_idx = gt_idx.to(device)
+        cls_targets[b, pred_idx] = 1
+
+        matched_coeffs = curve_coeffs[b, pred_idx]                # (M, 4)
+        pred_xs = eval_curve(matched_coeffs, sample_ys.unsqueeze(0).expand(pred_idx.numel(), -1))  # (M, S)
+        gt_xs = target_xs[b, gt_idx]                               # (M, S)
+        gt_mask = target_valid_mask[b, gt_idx]                     # (M, S)
+
+        if gt_mask.any():
+            reg_loss = F.smooth_l1_loss(pred_xs[gt_mask], gt_xs[gt_mask], reduction="mean")
+            reg_losses.append(reg_loss)
+
+        pred_endpoints = endpoints[b, pred_idx]                    # (M, 2)
+        gt_endpoints = target_endpoints[b, gt_idx]                 # (M, 2)
+        endpoint_losses.append(F.smooth_l1_loss(pred_endpoints, gt_endpoints, reduction="mean"))
+
+    class_weights = torch.tensor([bg_class_weight, 1.0], device=device)
+    cls_loss = F.cross_entropy(cls_logits.reshape(-1, 2), cls_targets.reshape(-1), weight=class_weights)
+
+    reg_loss = torch.stack(reg_losses).mean() if reg_losses else torch.zeros((), device=device)
+    endpoint_loss = torch.stack(endpoint_losses).mean() if endpoint_losses else torch.zeros((), device=device)
+
+    total = cls_weight * cls_loss + reg_weight * reg_loss + endpoint_weight * endpoint_loss
+
+    return {
+        "total": total,
+        "cls_loss": cls_loss.detach(),
+        "reg_loss": reg_loss.detach(),
+        "endpoint_loss": endpoint_loss.detach(),
+        "num_matched": sum(p.numel() for p, _ in matches),
+    }