File size: 3,160 Bytes
e2f3b24 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """
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):
# nearest keeps label values intact; 255 (ignore) survives too
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
|