| """ |
| Optional: compute the semantic loss at a reduced resolution. |
| |
| Ultralytics' SemanticSegmentationLoss upsamples the model's stride-8 logits to |
| the FULL mask resolution before the CE/Dice terms: |
| |
| preds = F.interpolate(preds, size=masks.shape[1:], ...) # 64x64 -> 512x512 |
| |
| At batch 32 / imgsz 512 / 19 classes that materializes a 159M-element tensor, |
| softmaxes it, then boolean-index-gathers it (a data-dependent shape that forces |
| an MPS sync) -- twice, counting the aux head. Measured, the loss costs more than |
| the entire forward+backward of the network. |
| |
| This patch computes the loss at `loss_size` instead (default 1/2 mask res): |
| predictions are upsampled only to loss_size and masks are nearest-downsampled to |
| match. Cost falls ~quadratically with loss_size. |
| |
| Trade-off: less sub-grid boundary supervision, so thin classes (eyes, brows, |
| lips) can suffer. Computing loss below the mask resolution is standard practice |
| in many semseg frameworks, but it IS a quality/speed trade -- validate before |
| trusting it. loss_size=0 restores stock behaviour. |
| |
| Usage: |
| import loss_patch; loss_patch.apply(loss_size=256) |
| """ |
| import torch |
| import torch.nn.functional as F |
| from ultralytics.utils import loss as _loss |
|
|
| _orig_forward = _loss.SemanticSegmentationLoss.forward |
| _applied = False |
|
|
|
|
| def apply(loss_size: int = 256): |
| """Patch SemanticSegmentationLoss.forward to evaluate at `loss_size`.""" |
| global _applied |
| if loss_size <= 0: |
| if _applied: |
| _loss.SemanticSegmentationLoss.forward = _orig_forward |
| _applied = False |
| return False |
|
|
| def forward(self, preds, batch): |
| aux_logits = None |
| if isinstance(preds, tuple): |
| preds, aux_logits = preds |
| masks = batch["semantic_mask"].to(preds.device) |
|
|
| h, w = masks.shape[1:] |
| th, tw = min(loss_size, h), min(loss_size, w) |
| if (th, tw) != (h, w): |
| |
| masks = F.interpolate(masks.float().unsqueeze(1), size=(th, tw), |
| mode="nearest").squeeze(1).to(masks.dtype) |
|
|
| valid = masks.reshape(-1) != 255 |
| if preds.shape[2:] != (th, tw): |
| preds = F.interpolate(preds, size=(th, tw), mode="bilinear", align_corners=False) |
|
|
| ce_loss = self._ce_loss(preds, masks, valid) |
| dice_loss = self._dice_loss(preds, masks, valid) |
| total = ce_loss + dice_loss |
|
|
| aux_loss = torch.tensor(0.0, device=preds.device, dtype=ce_loss.dtype) |
| if aux_logits is not None: |
| if aux_logits.shape[2:] != (th, tw): |
| aux_logits = F.interpolate(aux_logits, size=(th, tw), mode="bilinear", |
| align_corners=False) |
| aux_loss = self._ce_loss(aux_logits, masks, valid) * 0.4 |
| total += aux_loss |
|
|
| return total * preds.shape[0], {"ce_loss": ce_loss.detach(), |
| "dice_loss": dice_loss.detach(), |
| "aux_loss": aux_loss.detach()} |
|
|
| _loss.SemanticSegmentationLoss.forward = forward |
| _applied = True |
| return True |
|
|
