一、lovasz softmax魔改
    原lovasz softmax需要提供的标签为索引值，形状为(n,h,w)，如果使用外部库的话计算metric等不方便。
    新实现如下：

    1. """
    2. Lovasz-Softmax and Jaccard hinge loss in PyTorch
    3. Maxim Berman 2018 ESAT-PSI KU Leuven (MIT License)
    4. """
    6. from __future__ import print_function, division
    8. import torch
    9. from torch.autograd import Variable
    10. import torch.nn.functional as F
    11. import numpy as np
    12. try:
    13.     from itertools import  ifilterfalse
    14. except ImportError: # py3k
    15.     from itertools import  filterfalse as ifilterfalse
    16. from torch.nn.modules.loss import _Loss
    18. def lovasz_grad(gt_sorted):
    19.     """
    20.     Computes gradient of the Lovasz extension w.r.t sorted errors
    21.     See Alg. 1 in paper
    22.     """
    23.     p = len(gt_sorted)
    24.     gts = gt_sorted.sum()
    25.     intersection = gts - gt_sorted.float().cumsum(0)
    26.     union = gts + (1 - gt_sorted).float().cumsum(0)
    27.     jaccard = 1. - intersection / union
    28.     if p > 1: # cover 1-pixel case
    29.         jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
    30.     return jaccard
    33. def iou_binary(preds, labels, EMPTY=1., ignore=None, per_image=True):
    34.     """
    35.     IoU for foreground class
    36.     binary: 1 foreground, 0 background
    37.     """
    38.     if not per_image:
    39.         preds, labels = (preds,), (labels,)
    40.     ious = []
    41.     for pred, label in zip(preds, labels):
    42.         intersection = ((label == 1) & (pred == 1)).sum()
    43.         union = ((label == 1) | ((pred == 1) & (label != ignore))).sum()
    44.         if not union:
    45.             iou = EMPTY
    46.         else:
    47.             iou = float(intersection) / float(union)
    48.         ious.append(iou)
    49.     iou = mean(ious)    # mean accross images if per_image
    50.     return 100 * iou
    53. def iou(preds, labels, C, EMPTY=1., ignore=None, per_image=False):
    54.     """
    55.     Array of IoU for each (non ignored) class
    56.     """
    57.     if not per_image:
    58.         preds, labels = (preds,), (labels,)
    59.     ious = []
    60.     for pred, label in zip(preds, labels):
    61.         iou = []    
    62.         for i in range(C):
    63.             if i != ignore: # The ignored label is sometimes among predicted classes (ENet - CityScapes)
    64.                 intersection = ((label == i) & (pred == i)).sum()
    65.                 union = ((label == i) | ((pred == i) & (label != ignore))).sum()
    66.                 if not union:
    67.                     iou.append(EMPTY)
    68.                 else:
    69.                     iou.append(float(intersection) / float(union))
    70.         ious.append(iou)
    71.     ious = [mean(iou) for iou in zip(*ious)] # mean accross images if per_image
    72.     return 100 * np.array(ious)
    75. # --------------------------- BINARY LOSSES ---------------------------
    80. def lovasz_hinge(logits, labels, per_image=True, ignore=None):
    81.     """
    82.     Binary Lovasz hinge loss
    83.       logits: [B, H, W] Variable, logits at each pixel (between -\infty and +\infty)
    84.       labels: [B, H, W] Tensor, binary ground truth masks (0 or 1)
    85.       per_image: compute the loss per image instead of per batch
    86.       ignore: void class id
    87.     """
    88.     if per_image:
    89.         loss = mean(lovasz_hinge_flat(*flatten_binary_scores(log.unsqueeze(0), lab.unsqueeze(0), ignore))
    90.                           for log, lab in zip(logits, labels))
    91.     else:
    92.         loss = lovasz_hinge_flat(*flatten_binary_scores(logits, labels, ignore))
    93.     return loss
    96. def symmetric_lovasz_hinge(outputs, targets):
    97.     return (lovasz_hinge(outputs, targets) +
    98.             lovasz_hinge(-outputs, 1 - targets)) / 2
    102. def lovasz_hinge_flat(logits, labels):
    103.     """
    104.     Binary Lovasz hinge loss
    105.       logits: [P] Variable, logits at each prediction (between -\infty and +\infty)
    106.       labels: [P] Tensor, binary ground truth labels (0 or 1)
    107.       ignore: label to ignore
    108.     """
    109.     if len(labels) == 0:
    110.         # only void pixels, the gradients should be 0
    111.         return logits.sum() * 0.
    112.     signs = 2. * labels.float() - 1.
    113.     errors = (1. - logits * Variable(signs))
    114.     errors_sorted, perm = torch.sort(errors, dim=0, descending=True)
    115.     perm = perm.data
    116.     gt_sorted = labels[perm]
    117.     grad = lovasz_grad(gt_sorted)
    118.     loss = torch.dot(F.relu(errors_sorted), Variable(grad))
    119.     return loss
    122. def flatten_binary_scores(scores, labels, ignore=None):
    123.     """
    124.     Flattens predictions in the batch (binary case)
    125.     Remove labels equal to 'ignore'
    126.     """
    127.     scores = scores.view(-1)
    128.     labels = labels.view(-1)
    129.     if ignore is None:
    130.         return scores, labels
    131.     valid = (labels != ignore)
    132.     vscores = scores[valid]
    133.     vlabels = labels[valid]
    134.     return vscores, vlabels
    137. class StableBCELoss(torch.nn.modules.Module):
    138.     def __init__(self):
    139.          super(StableBCELoss, self).__init__()
    140.     def forward(self, input, target):
    141.          neg_abs = - input.abs()
    142.          loss = input.clamp(min=0) - input * target + (1 + neg_abs.exp()).log()
    143.          return loss.mean()
    146. def binary_xloss(logits, labels, ignore=None):
    147.     """
    148.     Binary Cross entropy loss
    149.       logits: [B, H, W] Variable, logits at each pixel (between -\infty and +\infty)
    150.       labels: [B, H, W] Tensor, binary ground truth masks (0 or 1)
    151.       ignore: void class id
    152.     """
    153.     logits, labels = flatten_binary_scores(logits, labels, ignore)
    154.     loss = StableBCELoss()(logits, Variable(labels.float()))
    155.     return loss
    158. # --------------------------- MULTICLASS LOSSES ---------------------------
    161. def lovasz_softmax(probas, labels, classes='present', per_image=False, ignore=None):
    162.     """
    163.     Multi-class Lovasz-Softmax loss
    164.       probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1).
    165.               Interpreted as binary (sigmoid) output with outputs of size [B, H, W].
    166.       labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1)
    167.       classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
    168.       per_image: compute the loss per image instead of per batch
    169.       ignore: void class labels
    170.     """
    171.     if per_image:
    172.         loss = mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), classes=classes)
    173.                           for prob, lab in zip(probas, labels))
    174.     else:
    175.         loss = lovasz_softmax_flat(*flatten_probas(probas, labels, ignore), classes=classes)
    176.     return loss
    179. class lovasz_softmax_onehot(_Loss):
    180.     def __init__(self):
    181.         """
    182.         原lovasz softmax需要提供的标签为索引值（n,h,w），这里魔改成提供onehot编码，方便计算IoU等指标
    183.         y_pred:(n,c,h,w)，应为softmax(dim=1)的值
    184.         y_true:(n,c,h,w)，应为onehot编码
    185.         """
    186.         super().__init__()
    188.     def forward(self, y_pred, y_true):
    189.         y_true = torch.argmax(y_true,dim=1)
    190.         loss = lovasz_softmax(y_pred,y_true)
    191.         return loss
    194. def lovasz_softmax_flat(probas, labels, classes='present'):
    195.     """
    196.     Multi-class Lovasz-Softmax loss
    197.       probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1)
    198.       labels: [P] Tensor, ground truth labels (between 0 and C - 1)
    199.       classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
    200.     """
    201.     if probas.numel() == 0:
    202.         # only void pixels, the gradients should be 0
    203.         return probas * 0.
    204.     C = probas.size(1)
    205.     losses = []
    206.     class_to_sum = list(range(C)) if classes in ['all', 'present'] else classes
    207.     for c in class_to_sum:
    208.         fg = (labels == c).float() # foreground for class c
    209.         if (classes is 'present' and fg.sum() == 0):
    210.             continue
    211.         if C == 1:
    212.             if len(classes) > 1:
    213.                 raise ValueError('Sigmoid output possible only with 1 class')
    214.             class_pred = probas[:, 0]
    215.         else:
    216.             class_pred = probas[:, c]
    217.         errors = (Variable(fg) - class_pred).abs()
    218.         errors_sorted, perm = torch.sort(errors, 0, descending=True)
    219.         perm = perm.data
    220.         fg_sorted = fg[perm]
    221.         losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted))))
    222.     return mean(losses)
    225. def flatten_probas(probas, labels, ignore=None):
    226.     """
    227.     Flattens predictions in the batch
    228.     """
    229.     if probas.dim() == 3:
    230.         # assumes output of a sigmoid layer
    231.         B, H, W = probas.size()
    232.         probas = probas.view(B, 1, H, W)
    233.     B, C, H, W = probas.size()
    234.     probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)  # B * H * W, C = P, C
    235.     labels = labels.view(-1)
    236.     if ignore is None:
    237.         return probas, labels
    238.     valid = (labels != ignore)
    239.     vprobas = probas[valid.nonzero().squeeze()]
    240.     vlabels = labels[valid]
    241.     return vprobas, vlabels
    243. def xloss(logits, labels, ignore=None):
    244.     """
    245.     Cross entropy loss
    246.     """
    247.     return F.cross_entropy(logits, Variable(labels), ignore_index=255)
    250. # --------------------------- HELPER FUNCTIONS ---------------------------
    251. def isnan(x):
    252.     return x != x
    255. def mean(l, ignore_nan=False, empty=0):
    256.     """
    257.     nanmean compatible with generators.
    258.     """
    259.     l = iter(l)
    260.     if ignore_nan:
    261.         l = ifilterfalse(isnan, l)
    262.     try:
    263.         n = 1
    264.         acc = next(l)
    265.     except StopIteration:
    266.         if empty == 'raise':
    267.             raise ValueError('Empty mean')
    268.         return empty
    269.     for n, v in enumerate(l, 2):
    270.         acc += v
    271.     if n == 1:
    272.         return acc
    273.     return acc / n

    使用仅需import lovasz_softmax_onehot即可，注意这里的pred是softmax(dim=1)后的，label为onehot编码。

    criterion = L.lovasz_softmax_onehot()
loss = criterion(pred,label)