Repository: PaddlePaddle/PaddleSeg

Issue:
多类别分割的类不均衡问题处理_自定义loss function
之前使用tensorflow 实现的自定义如下,可以使用,可以克服多类严重不均衡问题:
    def loss(self, logits, annotation):

        logits = tf.reshape(logits, (-1,num_of_classes))
        labels = tf.reshape(annotation, (-1,num_of_classes))
        cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels)
        
        weights = [0.1]
        for jj in range(1,num_of_classes):
            labels_temp = labels[...,jj]        
            positive_count = tf.count_nonzero(labels_temp,dtype=tf.float32)
            total_num = tf.constant(IMG_SIZE * IMG_SIZE,dtype=tf.float32)
            negativa_count = tf.subtract(total_num, positive_count)
            alpha = tf.div(tf.multiply(negativa_count,1.0),(positive_count*gama+1))
            weights = weights + [alpha]  
            
        weights = tf.convert_to_tensor(weights)   
        weights = tf.to_float(weights)
        labels = tf.to_float(labels)
    
        weight_map = tf.reduce_sum(tf.multiply(labels, weights), -1)
        weight_map = tf.stop_gradient(weight_map, name='stop_gradient')
        weighted_cross_entropy = tf.multiply(weight_map, cross_entropy, name='apply_weights')     
        loss_wc = tf.reduce_mean(weighted_cross_entropy, name='loss')        

        return loss_wc

然后我想转成paddle版本,如下:
def weighted_softmax_cross_entropy(logit, label,  num_classes=2):
                                                                
    logit = fluid.layers.reshape(logit, [-1,num_classes])
    label = fluid.layers.reshape(label, [-1,1])   
    
    logit = fluid.layers.softmax(logit)

    one_hot_label_ori = fluid.one_hot(label, depth=logit.shape[-1])
    one_hot_label = fluid.layers.squeeze(one_hot_label_ori, axes=[-2])
  
    log_prob = fluid.layers.log(fluid.layers.clip(logit, min=1e-7, max =1.))
    ce_loss = -1 * log_prob * one_hot_label 
    cross_entropy = fluid.layers.reduce_sum(ce_loss)  

    weight = fluid.layers.assign(np.array([0.01,1.0,2.,3.,5.], dtype='float32')) # result3 = [[2.5, 2.5], [2.5, 2.5], [2.5, 2.5]]
    weight.stop_gradient = True     
    weights=weight
    weight_map = fluid.layers.reduce_sum(fluid.layers.elementwise_mul(one_hot_label_ori, weights))    
    weight_map.stop_gradient = True 
    weighted_cross_entropy = fluid.layers.elementwise_mul(weight_map, cross_entropy)    
    loss_wc = fluid.layers.reduce_mean(weighted_cross_entropy)        

    return loss_wc

不知道我哪里写错了,训练的loss 不减少,不收敛
想请大家给看一下,谢谢


Discussion:

Resolving patch:
diff --git a/pdseg/lovasz_losses.py b/pdseg/lovasz_losses.py
new file mode 100755
--- /dev/null
+++ b/pdseg/lovasz_losses.py
@@ -0,0 +1,205 @@
+# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Lovasz-Softmax and Jaccard hinge loss in PaddlePaddle"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import paddle.fluid as fluid
+import numpy as np
+
+
+def _cumsum(x):
+    y = np.array(x)
+    return np.cumsum(y, axis=0)
+
+
+def create_tmp_var(name, dtype, shape):
+    return fluid.default_main_program().current_block().create_var(
+        name=name, dtype=dtype, shape=shape)
+
+
+def lovasz_grad(gt_sorted):
+    """
+    Computes gradient of the Lovasz extension w.r.t sorted errors
+    See Alg. 1 in paper
+    """
+    gt_sorted = fluid.layers.squeeze(gt_sorted, axes=[1])
+    gts = fluid.layers.reduce_sum(gt_sorted)
+    len_gt = fluid.layers.shape(gt_sorted)
+
+    # Acceleration is achieved by reducing the number of calls to cumsum.
+    # This calculation method is equivalent to that of the original paper.
+    var_one = fluid.layers.fill_constant(shape=[1], value=1, dtype='int32')
+    range_ = fluid.layers.range(1, len_gt + var_one, 1, 'int32')
+    tmp_var = create_tmp_var(
+        name='tmp_var', dtype=gt_sorted.dtype, shape=gt_sorted.shape)
+    cumsum_ = fluid.layers.py_func(func=_cumsum, x=gt_sorted, out=tmp_var)
+    intersection = gts - cumsum_
+    union = intersection + range_
+
+    jaccard = 1.0 - intersection / union
+    jaccard0 = fluid.layers.slice(jaccard, axes=[0], starts=[0], ends=[1])
+    jaccard1 = fluid.layers.slice(jaccard, axes=[0], starts=[1], ends=[len_gt])
+    jaccard2 = fluid.layers.slice(jaccard, axes=[0], starts=[0], ends=[-1])
+    jaccard = fluid.layers.concat([jaccard0, jaccard1 - jaccard2], axis=0)
+    jaccard = fluid.layers.unsqueeze(jaccard, axes=[1])
+    return jaccard
+
+
+def lovasz_hinge(logits, labels, ignore=None):
+    """
+    Binary Lovasz hinge loss
+      logits: [N, C, H, W] Tensor, logits at each pixel (between -\infty and +\infty)
+      labels: [N, 1, H, W] Tensor, binary ground truth masks (0 or 1)
+      ignore: [N, 1, H, W] Tensor. Void class labels, ignore pixels which value=0
+    """
+    loss = lovasz_hinge_flat(*flatten_binary_scores(logits, labels, ignore))
+    return loss
+
+
+def lovasz_hinge_flat(logits, labels):
+    """
+    Binary Lovasz hinge loss
+      logits: [P] Tensor, logits at each prediction (between -\infty and +\infty)
+      labels: [P] Tensor, binary ground truth labels (0 or 1)
+    """
+    shape = fluid.layers.shape(logits)
+    y = fluid.layers.zeros_like(shape[0])
+
+    out_var = fluid.layers.create_tensor("float32")
+    with fluid.layers.control_flow.Switch() as switch:
+        with switch.case(fluid.layers.equal(shape[0], y)):
+            loss = fluid.layers.reduce_sum(logits) * 0.
+            fluid.layers.assign(input=loss, output=out_var)
+        with switch.case(fluid.layers.greater_than(shape[0], y)):
+            labelsf = fluid.layers.cast(labels, logits.dtype)
+            signs = labelsf * 2 - 1.
+            signs.stop_gradient = True
+            errors = 1.0 - fluid.layers.elementwise_mul(logits, signs)
+            errors_sorted, perm = fluid.layers.argsort(
+                errors, axis=0, descending=True)
+            errors_sorted.stop_gradient = False
+            gt_sorted = fluid.layers.gather(labelsf, perm)
+
+            grad = lovasz_grad(gt_sorted)
+            grad.stop_gradient = True
+            loss = fluid.layers.reduce_sum(
+                fluid.layers.relu(errors_sorted) * grad)
+            fluid.layers.assign(input=loss, output=out_var)
+    return out_var
+
+
+def flatten_binary_scores(scores, labels, ignore=None):
+    """
+    Flattens predictions in the batch (binary case)
+    Remove labels according to 'ignore'
+    """
+    scores = fluid.layers.reshape(scores, [-1, 1])
+    labels = fluid.layers.reshape(labels, [-1, 1])
+    labels.stop_gradient = True
+    if ignore is None:
+        return scores, labels
+    ignore = fluid.layers.cast(ignore, 'int32')
+    ignore_mask = fluid.layers.reshape(ignore, (-1, 1))
+    indexs = fluid.layers.where(ignore_mask == 1)
+    indexs.stop_gradient = True
+    vscores = fluid.layers.gather(scores, indexs[:, 0])
+    vlabels = fluid.layers.gather(labels, indexs[:, 0])
+    return vscores, vlabels
+
+
+def lovasz_softmax(probas, labels, classes='present', ignore=None):
+    """
+    Multi-class Lovasz-Softmax loss
+      probas: [N, C, H, W] Tensor, class probabilities at each prediction (between 0 and 1).
+      labels: [N, 1, H, W] Tensor, ground truth labels (between 0 and C - 1)
+      classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
+      ignore: [N, 1, H, W] Tensor. Void class labels, ignore pixels which value=0
+    """
+    vprobas, vlabels = flatten_probas(probas, labels, ignore)
+    loss = lovasz_softmax_flat(vprobas, vlabels, classes=classes)
+    return loss
+
+
+def lovasz_softmax_flat(probas, labels, classes='present'):
+    """
+    Multi-class Lovasz-Softmax loss
+      probas: [P, C] Tensor, class probabilities at each prediction (between 0 and 1)
+      labels: [P] Tensor, ground truth labels (between 0 and C - 1)
+      classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
+    """
+    C = probas.shape[1]
+    losses = []
+    present = []
+    classes_to_sum = list(range(C)) if classes in ['all', 'present'
+                                                   ] else classes
+    for c in classes_to_sum:
+        fg = fluid.layers.cast(labels == c, probas.dtype)
+        fg.stop_gradient = True
+        if classes == 'present':
+            present.append(
+                fluid.layers.cast(fluid.layers.reduce_sum(fg) > 0, "int64"))
+        if C == 1:
+            if len(classes_to_sum) > 1:
+                raise ValueError('Sigmoid output possible only with 1 class')
+            class_pred = probas[:, 0]
+        else:
+            class_pred = probas[:, c]
+        errors = fluid.layers.abs(fg - class_pred)
+        errors_sorted, perm = fluid.layers.argsort(
+            errors, axis=0, descending=True)
+        errors_sorted.stop_gradient = False
+
+        fg_sorted = fluid.layers.gather(fg, perm)
+        fg_sorted.stop_gradient = True
+
+        grad = lovasz_grad(fg_sorted)
+        grad.stop_gradient = True
+        loss = fluid.layers.reduce_sum(errors_sorted * grad)
+
+        losses.append(loss)
+
+    if len(classes_to_sum) == 1:
+        return losses[0]
+
+    losses_tensor = fluid.layers.stack(losses)
+    if classes == 'present':
+        present_tensor = fluid.layers.stack(present)
+        index = fluid.layers.where(present_tensor == 1)
+        index.stop_gradient = True
+        losses_tensor = fluid.layers.gather(losses_tensor, index[:, 0])
+    loss = fluid.layers.mean(losses_tensor)
+    return loss
+
+
+def flatten_probas(probas, labels, ignore=None):
+    """
+    Flattens predictions in the batch
+    """
+    if len(probas.shape) == 3:
+        probas = fluid.layers.unsqueeze(probas, axis=[1])
+    C = probas.shape[1]
+    probas = fluid.layers.transpose(probas, [0, 2, 3, 1])
+    probas = fluid.layers.reshape(probas, [-1, C])
+    labels = fluid.layers.reshape(labels, [-1, 1])
+    if ignore is None:
+        return probas, labels
+    ignore = fluid.layers.cast(ignore, 'int32')
+    ignore_mask = fluid.layers.reshape(ignore, [-1, 1])
+    indexs = fluid.layers.where(ignore_mask == 1)
+    indexs.stop_gradient = True
+    vprobas = fluid.layers.gather(probas, indexs[:, 0])
+    vlabels = fluid.layers.gather(labels, indexs[:, 0])
+    return vprobas, vlabels
diff --git a/pdseg/models/model_builder.py b/pdseg/models/model_builder.py
--- a/pdseg/models/model_builder.py
+++ b/pdseg/models/model_builder.py
@@ -24,6 +24,8 @@
 from loss import multi_softmax_with_loss
 from loss import multi_dice_loss
 from loss import multi_bce_loss
+from lovasz_losses import lovasz_hinge
+from lovasz_losses import lovasz_softmax
 from models.modeling import deeplab, unet, icnet, pspnet, hrnet, fast_scnn
 
 
@@ -204,19 +206,22 @@ def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
             if not isinstance(loss_type, list):
                 loss_type = list(loss_type)
 
-            # dice_loss或bce_loss只适用两类分割中
-            if class_num > 2 and (("dice_loss" in loss_type) or
+            # lovasz_hinge_loss或dice_loss或bce_loss只适用两类分割中
+            if class_num > 2 and (("lovasz_hinge_loss" in loss_type) or
+                                  ("dice_loss" in loss_type) or
                                   ("bce_loss" in loss_type)):
                 raise Exception(
-                    "dice loss and bce loss is only applicable to binary classfication"
+                    "lovasz hinge loss, dice loss and bce loss are only applicable to binary classfication."
                 )
 
-            # 在两类分割情况下，当loss函数选择dice_loss或bce_loss的时候，最后logit输出通道数设置为1
-            if ("dice_loss" in loss_type) or ("bce_loss" in loss_type):
+            # 在两类分割情况下，当loss函数选择lovasz_hinge_loss或dice_loss或bce_loss的时候，最后logit输出通道数设置为1
+            if ("dice_loss" in loss_type) or ("bce_loss" in loss_type) or (
+                    "lovasz_hinge_loss" in loss_type):
                 class_num = 1
-                if "softmax_loss" in loss_type:
+                if ("softmax_loss" in loss_type) or (
+                        "lovasz_softmax_loss" in loss_type):
                     raise Exception(
-                        "softmax loss can not combine with dice loss or bce loss"
+                        "softmax loss or lovasz softmax loss can not combine with bce loss or dice loss or lovasz hinge loss."
                     )
             logits = seg_model(image, class_num)
 
@@ -240,11 +245,22 @@ def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
                     avg_loss_list.append(multi_bce_loss(logits, label, mask))
                     loss_valid = True
                     valid_loss.append("bce_loss")
+                if "lovasz_hinge_loss" in loss_type:
+                    avg_loss_list.append(
+                        lovasz_hinge(logits, label, ignore=mask))
+                    loss_valid = True
+                    valid_loss.append("lovasz_hinge_loss")
+                if "lovasz_softmax_loss" in loss_type:
+                    probas = fluid.layers.softmax(logits, axis=1)
+                    avg_loss_list.append(
+                        lovasz_softmax(probas, label, ignore=mask))
+                    loss_valid = True
+                    valid_loss.append("lovasz_softmax_loss")
                 if not loss_valid:
                     raise Exception(
                         "SOLVER.LOSS: {} is set wrong. it should "
-                        "include one of (softmax_loss, bce_loss, dice_loss) at least"
-                        " example: ['softmax_loss'], ['dice_loss'], ['bce_loss', 'dice_loss']"
+                        "include one of (softmax_loss, bce_loss, dice_loss, lovasz_hinge_loss, lovasz_softmax_loss) at least"
+                        " example: ['softmax_loss'], ['dice_loss'], ['bce_loss', 'dice_loss'], ['lovasz_hinge_loss','bce_loss'], ['lovasz_softmax_loss','softmax_loss']"
                         .format(cfg.SOLVER.LOSS))
 
                 invalid_loss = [x for x in loss_type if x not in valid_loss]
@@ -255,7 +271,9 @@ def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
 
                 avg_loss = 0
                 for i in range(0, len(avg_loss_list)):
-                    avg_loss += avg_loss_list[i]
+                    loss_name = valid_loss[i].upper()
+                    loss_weight = eval('cfg.SOLVER.LOSS_WEIGHT.' + loss_name)
+                    avg_loss += loss_weight * avg_loss_list[i]
 
             #get pred result in original size
             if isinstance(logits, tuple):
@@ -268,7 +286,7 @@ def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
 
             # return image input and logit output for inference graph prune
             if ModelPhase.is_predict(phase):
-                # 两类分割中，使用dice_loss或bce_loss返回的logit为单通道，进行到两通道的变换
+                # 两类分割中，使用lovasz_hinge_loss或dice_loss或bce_loss返回的logit为单通道，进行到两通道的变换
                 if class_num == 1:
                     logit = sigmoid_to_softmax(logit)
                 else:
diff --git a/pdseg/utils/config.py b/pdseg/utils/config.py
--- a/pdseg/utils/config.py
+++ b/pdseg/utils/config.py
@@ -155,6 +155,12 @@
 cfg.SOLVER.NUM_EPOCHS = 30
 # loss的选择，支持softmax_loss, bce_loss, dice_loss
 cfg.SOLVER.LOSS = ["softmax_loss"]
+# loss的权重，用于多loss组合加权使用，仅对SOLVER.LOSS内包含的loss生效
+cfg.SOLVER.LOSS_WEIGHT.SOFTMAX_LOSS = 1
+cfg.SOLVER.LOSS_WEIGHT.DICE_LOSS = 1
+cfg.SOLVER.LOSS_WEIGHT.BCE_LOSS = 1
+cfg.SOLVER.LOSS_WEIGHT.LOVASZ_HINGE_LOSS = 1
+cfg.SOLVER.LOSS_WEIGHT.LOVASZ_SOFTMAX_LOSS = 1
 # 是否开启warmup学习策略
 cfg.SOLVER.LR_WARMUP = False
 # warmup的迭代次数


Test patch: