| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
|
|
|
|
| class CIBCELoss(nn.Module): |
| def __init__( |
| self, |
| output_key: str = "ci", |
| label_key: str = "ci", |
| weight: float = 1.0, |
| eps: float = 1e-6, |
| ): |
| super().__init__() |
| self.output_key = str(output_key) |
| self.label_key = str(label_key) |
| self.weight = float(weight) |
| self.eps = float(eps) |
|
|
| def forward(self, outputs: dict[str, torch.Tensor], labels: dict[str, torch.Tensor]) -> torch.Tensor: |
| preds = torch.clamp(outputs[self.output_key].float(), self.eps, 1.0 - self.eps) |
| targets = labels[self.label_key].float() |
| if targets.ndim == 4 and targets.shape[1] == 1: |
| targets = targets[:, 0] |
| loss = self.weight * F.binary_cross_entropy(preds, targets) |
| self.last_components = {"bce": float(loss.detach().cpu())} |
| return loss |
|
|
|
|
| class CIFocalLoss(nn.Module): |
| def __init__( |
| self, |
| output_key: str = "ci", |
| label_key: str = "ci", |
| weight: float = 1.0, |
| alpha: float = 0.25, |
| gamma: float = 2.0, |
| eps: float = 1e-6, |
| ): |
| super().__init__() |
| self.output_key = str(output_key) |
| self.label_key = str(label_key) |
| self.weight = float(weight) |
| self.alpha = float(alpha) |
| self.gamma = float(gamma) |
| self.eps = float(eps) |
|
|
| def forward(self, outputs: dict[str, torch.Tensor], labels: dict[str, torch.Tensor]) -> torch.Tensor: |
| preds = torch.clamp(outputs[self.output_key].float(), self.eps, 1.0 - self.eps) |
| targets = labels[self.label_key].float().clamp(0.0, 1.0) |
| if targets.ndim == 4 and targets.shape[1] == 1: |
| targets = targets[:, 0] |
|
|
| bce = F.binary_cross_entropy(preds, targets, reduction="none") |
| pt = targets * preds + (1.0 - targets) * (1.0 - preds) |
| alpha_t = targets * self.alpha + (1.0 - targets) * (1.0 - self.alpha) |
| loss = self.weight * (alpha_t * (1.0 - pt).pow(self.gamma) * bce).mean() |
| self.last_components = {"focal": float(loss.detach().cpu())} |
| return loss |
|
|
|
|
| class CIFocalTverskyLoss(nn.Module): |
| def __init__( |
| self, |
| output_key: str = "ci", |
| label_key: str = "ci", |
| weight: float = 1.0, |
| focal_weight: float = 0.5, |
| tversky_weight: float = 0.5, |
| fp_weight: float = 0.7, |
| fn_weight: float = 0.3, |
| focal_gamma: float = 2.0, |
| tversky_gamma: float = 1.0, |
| positive_alpha: float = 0.35, |
| target_threshold: float = 0.5, |
| eps: float = 1e-6, |
| ): |
| super().__init__() |
| self.output_key = str(output_key) |
| self.label_key = str(label_key) |
| self.weight = float(weight) |
| self.focal_weight = float(focal_weight) |
| self.tversky_weight = float(tversky_weight) |
| self.fp_weight = float(fp_weight) |
| self.fn_weight = float(fn_weight) |
| self.focal_gamma = float(focal_gamma) |
| self.tversky_gamma = float(tversky_gamma) |
| self.positive_alpha = float(positive_alpha) |
| self.target_threshold = float(target_threshold) |
| self.eps = float(eps) |
|
|
| def forward(self, outputs: dict[str, torch.Tensor], labels: dict[str, torch.Tensor]) -> torch.Tensor: |
| probs = torch.clamp(outputs[self.output_key].float(), self.eps, 1.0 - self.eps) |
| targets = labels[self.label_key].float().clamp(0.0, 1.0) |
| if targets.ndim == 4 and targets.shape[1] == 1: |
| targets = targets[:, 0] |
|
|
| bce = F.binary_cross_entropy(probs, targets, reduction="none") |
| pt = targets * probs + (1.0 - targets) * (1.0 - probs) |
| alpha_t = targets * self.positive_alpha + (1.0 - targets) * (1.0 - self.positive_alpha) |
| focal = alpha_t * (1.0 - pt).pow(self.focal_gamma) * bce |
| focal_loss = focal.mean() |
|
|
| reduce_dims = tuple(range(1, probs.dim())) |
| y_hard = (targets >= self.target_threshold).float() |
| tp = (probs * y_hard).sum(dim=reduce_dims) |
| fp = (probs * (1.0 - y_hard)).sum(dim=reduce_dims) |
| fn = ((1.0 - probs) * y_hard).sum(dim=reduce_dims) |
| tversky = (tp + self.eps) / (tp + self.fp_weight * fp + self.fn_weight * fn + self.eps) |
| tversky_loss = (1.0 - tversky).pow(self.tversky_gamma).mean() |
|
|
| loss = self.focal_weight * focal_loss + self.tversky_weight * tversky_loss |
| total = self.weight * loss |
| self.last_components = { |
| "focal": float((self.weight * self.focal_weight * focal_loss).detach().cpu()), |
| "tversky": float((self.weight * self.tversky_weight * tversky_loss).detach().cpu()), |
| "total": float(total.detach().cpu()), |
| } |
| return total |
|
|
|
|
| class BTMaskedSmoothL1Loss(nn.Module): |
| def __init__( |
| self, |
| output_key: str = "bt", |
| label_key: str = "bt", |
| weight: float = 1.0, |
| lead_weights: list[float] | None = None, |
| beta: float = 1.0, |
| ): |
| super().__init__() |
| self.output_key = str(output_key) |
| self.label_key = str(label_key) |
| self.weight = float(weight) |
| self.beta = float(beta) |
| self.register_buffer("lead_weights", torch.tensor(lead_weights or [], dtype=torch.float32), persistent=False) |
|
|
| def forward(self, outputs: dict[str, torch.Tensor], labels: dict[str, torch.Tensor]) -> torch.Tensor: |
| preds = outputs[self.output_key].float() |
| targets = labels[self.label_key].float() |
| valid = torch.isfinite(targets) |
| if not bool(valid.any()): |
| zero = preds.sum() * 0.0 |
| self.last_components = {"bt_smooth_l1": 0.0} |
| return zero |
|
|
| loss = F.smooth_l1_loss(preds, torch.nan_to_num(targets), beta=self.beta, reduction="none") |
| if self.lead_weights.numel() > 0: |
| if self.lead_weights.numel() != preds.shape[1]: |
| raise ValueError(f"lead_weights length {self.lead_weights.numel()} != T {preds.shape[1]}") |
| view_shape = (1, preds.shape[1]) + (1,) * (preds.ndim - 2) |
| loss = loss * self.lead_weights.to(preds.device).view(view_shape) |
| loss = self.weight * loss[valid].mean() |
| self.last_components = {"bt_smooth_l1": float(loss.detach().cpu())} |
| return loss |
|
|
|
|
| class BTMaskedMSELoss(nn.Module): |
| def __init__( |
| self, |
| output_key: str = "bt", |
| label_key: str = "bt", |
| weight: float = 1.0, |
| lead_weights: list[float] | None = None, |
| ): |
| super().__init__() |
| self.output_key = str(output_key) |
| self.label_key = str(label_key) |
| self.weight = float(weight) |
| self.register_buffer("lead_weights", torch.tensor(lead_weights or [], dtype=torch.float32), persistent=False) |
|
|
| def forward(self, outputs: dict[str, torch.Tensor], labels: dict[str, torch.Tensor]) -> torch.Tensor: |
| preds = outputs[self.output_key].float() |
| targets = labels[self.label_key].float() |
| valid = torch.isfinite(targets) |
| if not bool(valid.any()): |
| zero = preds.sum() * 0.0 |
| self.last_components = {"bt_mse": 0.0} |
| return zero |
|
|
| loss = (preds - torch.nan_to_num(targets)) ** 2 |
| if self.lead_weights.numel() > 0: |
| if self.lead_weights.numel() != preds.shape[1]: |
| raise ValueError(f"lead_weights length {self.lead_weights.numel()} != T {preds.shape[1]}") |
| view_shape = (1, preds.shape[1]) + (1,) * (preds.ndim - 2) |
| loss = loss * self.lead_weights.to(preds.device).view(view_shape) |
| loss = self.weight * loss[valid].mean() |
| self.last_components = {"bt_mse": float(loss.detach().cpu())} |
| return loss |
|
|
|
|
| class CIBTLoss(nn.Module): |
| def __init__( |
| self, |
| ci_loss: dict, |
| bt_loss: dict | None = None, |
| normalize_terms: bool = False, |
| ref_batches: int = 100, |
| eps: float = 1e-8, |
| ): |
| super().__init__() |
| ci_cfg = dict(ci_loss) |
| bt_cfg = dict(bt_loss or {"name": "bt_mse", "weight": 0.05}) |
| self.ci_weight = float(ci_cfg.pop("weight", 1.0)) |
| self.bt_weight = float(bt_cfg.pop("weight", 1.0)) |
| self.ci_loss = _build_one_loss(ci_cfg, default_name=str(ci_cfg.get("name", "bce"))) |
| self.bt_loss = _build_one_loss(bt_cfg, default_name="bt_mse") |
| self.normalize_terms = bool(normalize_terms) |
| self.ref_batches = int(ref_batches) |
| self.eps = float(eps) |
| self.register_buffer("ref_count", torch.tensor(0.0), persistent=True) |
| self.register_buffer("ci_ref_sum", torch.tensor(0.0), persistent=True) |
| self.register_buffer("bt_ref_sum", torch.tensor(0.0), persistent=True) |
| self.register_buffer("ci_ref", torch.tensor(1.0), persistent=True) |
| self.register_buffer("bt_ref", torch.tensor(1.0), persistent=True) |
| self.last_components: dict[str, float] = {} |
|
|
| def forward(self, outputs: dict[str, torch.Tensor], labels: dict[str, torch.Tensor]) -> torch.Tensor: |
| ci = self.ci_loss(outputs, labels) |
| bt = self.bt_loss(outputs, labels) |
| if self.normalize_terms: |
| self._update_refs(ci, bt) |
| ci_term = ci / self.ci_ref.clamp_min(self.eps) |
| bt_term = bt / self.bt_ref.clamp_min(self.eps) |
| else: |
| ci_term = ci |
| bt_term = bt |
| ci_weighted = self.ci_weight * ci_term |
| bt_weighted = self.bt_weight * bt_term |
| total = ci_weighted + bt_weighted |
|
|
| components = { |
| "ci": float(ci_weighted.detach().cpu()), |
| "bt": float(bt_weighted.detach().cpu()), |
| "total": float(total.detach().cpu()), |
| "ci_raw": float(ci.detach().cpu()), |
| "bt_raw": float(bt.detach().cpu()), |
| "ci_norm": float(ci_term.detach().cpu()), |
| "bt_norm": float(bt_term.detach().cpu()), |
| "ci_weight": float(self.ci_weight), |
| "bt_weight": float(self.bt_weight), |
| "ci_ref": float(self.ci_ref.detach().cpu()), |
| "bt_ref": float(self.bt_ref.detach().cpu()), |
| "ref_count": float(self.ref_count.detach().cpu()), |
| "ref_ready": float(self._ref_ready()), |
| } |
| for name, value in getattr(self.ci_loss, "last_components", {}).items(): |
| components[f"ci_{name}"] = float(value) |
| for name, value in getattr(self.bt_loss, "last_components", {}).items(): |
| component_name = str(name) if str(name).startswith("bt_") else f"bt_{name}" |
| components[component_name] = float(value) |
| self.last_components = components |
| return total |
|
|
| def _ref_ready(self) -> bool: |
| return bool(self.normalize_terms and self.ref_count.item() >= max(self.ref_batches, 1)) |
|
|
| def _update_refs(self, ci: torch.Tensor, bt: torch.Tensor) -> None: |
| if self.ref_batches <= 0: |
| if self.ref_count.item() < 1: |
| self.ci_ref.copy_(ci.detach().abs().clamp_min(self.eps)) |
| self.bt_ref.copy_(bt.detach().abs().clamp_min(self.eps)) |
| self.ref_count.fill_(1.0) |
| return |
| if self.ref_count.item() >= self.ref_batches: |
| return |
| self.ci_ref_sum.add_(ci.detach().abs()) |
| self.bt_ref_sum.add_(bt.detach().abs()) |
| self.ref_count.add_(1.0) |
| denom = self.ref_count.clamp_min(1.0) |
| self.ci_ref.copy_((self.ci_ref_sum / denom).clamp_min(self.eps)) |
| self.bt_ref.copy_((self.bt_ref_sum / denom).clamp_min(self.eps)) |
|
|
|
|
| class CombinedLoss(nn.Module): |
| def __init__(self, losses: dict[str, nn.Module]): |
| super().__init__() |
| if not losses: |
| raise ValueError("CombinedLoss requires at least one loss") |
| self.losses = nn.ModuleDict(losses) |
| self.last_components: dict[str, float] = {} |
|
|
| def forward(self, outputs: dict[str, torch.Tensor], labels: dict[str, torch.Tensor]) -> torch.Tensor: |
| total = None |
| components = {} |
| for name, loss_fn in self.losses.items(): |
| value = loss_fn(outputs, labels) |
| components[name] = float(value.detach().cpu()) |
| for sub_name, sub_value in getattr(loss_fn, "last_components", {}).items(): |
| component_name = str(sub_name) if name == "loss" else f"{name}_{sub_name}" |
| components[component_name] = float(sub_value) |
| total = value if total is None else total + value |
| self.last_components = components |
| if total is None: |
| raise RuntimeError("no losses were evaluated") |
| return total |
|
|
|
|
| def build_loss(config: dict) -> nn.Module: |
| config = dict(config) |
| if str(config.get("name", "")).lower() == "ci_bt": |
| return _build_one_loss(config, default_name="ci_bt") |
| if "losses" in config: |
| return CombinedLoss( |
| { |
| str(component_name): _build_one_loss(dict(component_cfg), default_name=str(component_name)) |
| for component_name, component_cfg in dict(config["losses"]).items() |
| } |
| ) |
| return CombinedLoss({"loss": _build_one_loss(config, default_name=str(config.get("name", "binary_focal")))}) |
|
|
|
|
| def _build_one_loss(config: dict, default_name: str) -> nn.Module: |
| name = str(config.get("name", default_name)).lower() |
| params = dict(config.get("params", {})) |
| output_key = str(config.get("output_key", "bt" if name.startswith("bt") or "smooth_l1" in name else "ci")) |
| label_key = str(config.get("label_key", config.get("target_label", output_key))) |
| weight = float(config.get("weight", 1.0)) |
|
|
| if name in {"bce", "binary_bce", "binary_cross_entropy", "ci_bce"}: |
| return CIBCELoss(output_key=output_key, label_key=label_key, weight=weight, **params) |
| if name in {"binary_focal", "focal", "binary_focal_loss", "ci", "ci_focal"}: |
| return CIFocalLoss(output_key=output_key, label_key=label_key, weight=weight, **params) |
| if name in {"focal_tversky", "tversky_focal", "ci_focal_tversky", "far_aware_tversky_focal", "far_aware", "ci_far_aware"}: |
| return CIFocalTverskyLoss(output_key=output_key, label_key=label_key, weight=weight, **params) |
| if name in {"masked_lead_smooth_l1", "bt_smooth_l1", "bt", "bt_masked_smooth_l1"}: |
| return BTMaskedSmoothL1Loss(output_key=output_key, label_key=label_key, weight=weight, **params) |
| if name in {"bt_mse", "masked_bt_mse", "bt_masked_mse"}: |
| return BTMaskedMSELoss(output_key=output_key, label_key=label_key, weight=weight, **params) |
| if name == "ci_bt": |
| if "ci_loss" not in config: |
| raise ValueError("ci_bt loss requires a ci_loss config block") |
| bt_loss = config.get("bt_loss") |
| if bt_loss is None: |
| bt_loss = { |
| "name": "bt_mse", |
| "output_key": config.get("bt_output_key", "bt"), |
| "label_key": config.get("bt_label_key", "bt"), |
| "weight": config.get("bt_weight", 0.05), |
| } |
| return CIBTLoss( |
| ci_loss=dict(config["ci_loss"]), |
| bt_loss=dict(bt_loss), |
| **params, |
| ) |
| raise ValueError(f"unknown loss: {name}") |
|
|
