File size: 2,968 Bytes
71d64bb | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """
Lovász-Softmax loss — NaN-safe version for small-batch training.
Fixes over v1:
- eps guard on union denominator (prevents /0 on rare classes)
- support for multiple ignore indices (background + void)
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
def lovasz_grad(gt_sorted):
"""Gradient of Lovász extension w.r.t. sorted errors."""
p = len(gt_sorted)
gts = gt_sorted.sum()
intersection = gts - gt_sorted.float().cumsum(0)
union = gts + (1.0 - gt_sorted).float().cumsum(0)
jaccard = 1.0 - intersection / (union + 1e-6) # eps guard
if p > 1:
jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
return jaccard
def lovasz_softmax_flat(probas, labels, classes="present"):
"""
Multi-class Lovász-Softmax on flattened tensors.
probas: (N, C) softmax probabilities
labels: (N,) ground truth indices
"""
if probas.numel() == 0:
return probas * 0.0
C = probas.shape[1]
if classes == "all":
class_list = list(range(C))
elif classes == "present":
class_list = torch.unique(labels).tolist()
else:
class_list = list(classes)
losses = []
for c in class_list:
fg = (labels == c).float()
if fg.sum() == 0:
continue
class_pred = probas[:, c]
errors = (fg - class_pred).abs()
errors_sorted, perm = torch.sort(errors, dim=0, descending=True)
fg_sorted = fg[perm.data]
grad = lovasz_grad(fg_sorted)
loss_c = torch.dot(errors_sorted, grad)
if torch.isfinite(loss_c):
losses.append(loss_c)
if not losses:
return probas.sum() * 0.0
return sum(losses) / len(losses)
class LovaszSoftmaxLoss(nn.Module):
"""
Lovász-Softmax for 2D segmentation.
logits: (B, C, H, W)
targets: (B, H, W) long
ignore_indices: list of class indices to exclude (e.g. [0, 19])
"""
def __init__(self, ignore_indices=None, classes="present"):
super().__init__()
self.ignore_indices = ignore_indices or []
self.classes = classes
def forward(self, logits, targets):
probas = F.softmax(logits, dim=1)
B, C, H, W = probas.shape
probas_flat = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)
targets_flat = targets.contiguous().view(-1)
# Mask out all ignore indices
if self.ignore_indices:
valid = torch.ones(targets_flat.shape[0], dtype=torch.bool,
device=targets_flat.device)
for idx in self.ignore_indices:
valid &= (targets_flat != idx)
probas_flat = probas_flat[valid]
targets_flat = targets_flat[valid]
if probas_flat.numel() == 0:
return logits.sum() * 0.0
loss = lovasz_softmax_flat(probas_flat, targets_flat, classes=self.classes)
return torch.nan_to_num(loss, nan=0.0)
|