医疗图像分割的损失函数

写在前面

图像分割是一个很基础的计算机视觉的问题，最近在我的研究方向中遇到的图像分割问题，就查阅了一些文献。由于我的项目主要用到的MRI图像，就自然而然的使用到了Unet¹。使用的loss也是最原始的Dice系数。由于loss是需要一直减小的，而Dice系数则接近1越好。因此这里的一个trick就是将loss设置为_-Dice_。在实验的过程中，发现loss下降到-0.6之后就不再下降了。这是一个很差的结果，然后我将Validation set 的数据predict出segmentation Map，结果自然而然不好，大概能够给出要分割区域的位置，然而实际上边缘十分的粗糙，方方正正的（我想不明白这里面有什么玄学）。最终我决定从loss 函数入手去解决这个问题。因为我发现很多文章中提出对于分割来说，由于最终要得到的是segmentation Map，因此相当于就是对于每一个像素点的predictioin。这样的话一个重要的方面就是我们要分割的区域通常是占据小部分面积，这样就造成训练样例间的不均衡（这当然是有待考究的）。我想先从基本的Unet模型下手，通过尝试不同的loss function。看看实验结果会不会有什么改进。当然模型也是一个很重要的方面，我想接下来肯定要读论文进行一些总结的。
在CSDN上有许多很好的参考资源，而我参考就是其中一篇。从loss方面处理图像分割中数据极度不均衡的状况 --keras。²总结的很全面，我想我接下来的一些实验也会根据这篇Blog里面描述的一些内容并根据我的项目场景进行。³

Tversky_loss

"""
# Ref: salehi17, "Twersky loss function for image segmentation using 3D FCDN"
# -> the score is computed for each class separately and then summed
# alpha=beta=0.5 : dice coefficient
# alpha=beta=1   : tanimoto coefficient (also known as jaccard)
# alpha+beta=1   : produces set of F*-scores
# implemented by E. Moebel, 06/04/18
"""
def tversky_loss(y_true, y_pred):
    alpha = 0.5
    beta  = 0.5
    
    ones = K.ones(K.shape(y_true))
    p0 = y_pred      # proba that voxels are class i
    p1 = ones-y_pred # proba that voxels are not class i
    g0 = y_true
    g1 = ones-y_true
    
    num = K.sum(p0*g0, (0,1,2))
    den = num + alpha*K.sum(p0*g1,(0,1,2)) + beta*K.sum(p1*g0,(0,1,2))
    
    T = K.sum(num/den) # when summing over classes, T has dynamic range [0 Ncl]
    
    Ncl = K.cast(K.shape(y_true)[-1], 'float32') 
    return Ncl-T

这里是针对多类分割问题的loss，也可以用于binary的分割
这里面有两个超参数 alpha, beta.
alpha = beta =0.5 那么就是Dice loss
alpha + beta =1 那么就是Dice loss的一些variant, 取决于到底是precision还是sensitity 重要。
alpha = beta =1 就是IOU loss或者jaccard loss, 一般来说， 同样的GT和PR，IOU loss 要比 Dice loss大，对于我们来说，loss越大越好，因为更容易寻找梯度下降的方向，换言之就是更容易优化。那么这只能表明IOU 比 Dice 收敛快一些，至于最终的收敛值，就因问题而异，实践一下。
同理 alpha+beta = 2是对权重的一个调整，可据问题选择。

对DICE系数的思考

通常我更换损失函数时，会使用相同的metric评估模型,可是更换了模型和损失函数之后，dice系数没有按照我理想的方向改进，然后我就重新思考了一下dice系数的公式。
$$Dice(A,B)=2\ast \frac{A\cap B}{A+B}$$
虽然Dice系数值可能相差不多，但是可能会是两种不同的情况。我们先假设A是Ground Truth,以下简称GT, B是Prediction，以下简称PR。

1. B可能比较小，也就是PR小于GT.
2. B比较大，也就是PR的区域是要大于GT。
   这两种情况下，可能会出现相同的分数。在医疗应用中，recall通常是比较重要的，所以我们会倾向于第二种情况，如果他们的分数相差不多的情况下。那么进一步怎么去比较到底是那种情况呢，这里我们就引出新的损失函数，IOU.

对IOU的一些思考

IOU通常在object detection里面用的比较多，用来评测PR的boundding box和GT之间的误差。
$$IOU(A,B)=\frac{A\cap B}{A\cup B}$$
我们可以得到$IOU(A,B)\in [0,1]$,这个性质很好，可以判断我们的函数是否工作，并且在training的过程中，损失值递减。
那么刚才说Dice系数值可能相同的两种情况，我们就可以用IOU作为第二评测指标。⁴

Focal loss

Focal loss 来源于FBIR的论文，原意是为了解决物体检测中的ROI背景与前景不平衡问题，我想对于我们的问题同样适用，只不过是从像素level.
代码：https://github.com/mkocabas/focal-loss-keras

from keras import backend as K 
"""Compatible with tensorflow backend"""
def focal_loss(gamma=2., alpha=.25):
	def focal_loss_fixed(y_true, y_pred):
		pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
        pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
        return -K.sum(alpha * K.pow(1. - pt_1, gamma) * K.log(pt_1))-K.sum((1-alpha) * K.pow( pt_0, gamma) * K.log(1. - pt_0))
	return focal_loss_fixed

先做实验试一下
容易发生溢出,还没搞清为什么

Generalizeddiceloss with class weights

针对多种类的数据不平衡问题，当然可以同样尝试使用给loss添加权重去解决。权重可以是基于数据集提前计算出来的，当然也可以是根据batch数据计算出来的，当然如果batch size 太小，会产生一定的震荡。

def generalized_dice_loss_w(y_true, y_pred): 
    # Compute weights: "the contribution of each label is corrected by the inverse of its volume"
    Ncl = y_pred.shape[-1]
    w = K.zeros((Ncl,))
    # w = K.sum(y_true, (0,1,2))
    # w = K.sum(w)/(w + 1.0)
    # w = w/K.sum(w)
    w = np.array([0, 0.20,  0.29,  0.01,  0.5])*(1-alpha)

    numerator = y_true*y_pred
    numerator = w*K.sum(numerator, (0,1,2))
    denominator = y_true+y_pred+1e-5
    denominator = K.sum(denominator, (0,1,2))
    gen_dice_coef = numerator/denominator
    gen_dice_coef = K.sum(gen_dice_coef)

    return 1-2*gen_dice_coef

BCE（Binary crossentropy）

截止目前能够达到最好的效果，dice=91.36
使用预训练的权重，可以提高1.5个百分点。
本来我以为Unet只能做binary的segmenation,这样就太鸡肋了，后来才发现原来是我没理解到位。

binary_crossentropy(y_true, y_pred)

对于多类分割就换成multiclass的。同样可能遭遇的问题就是待分割种类数据不均衡。

总结

损失函数的设计是用来处理数据集具有的一些特殊的情况，比方说Focal loss就是用来处理数据集中的难分样本。DICE系数可以用来处理数据分类不均衡的情况，其中不均衡很多情况下是由于背景和待分割区域之间的面积对比不均衡。对于Binary的classification, 我们当然可以只考虑待分割区域，即是本文中的DICE函数。那么当对于多种类的分割时，我们同样可以只计算待分割区域的DICE系数，这样就可以避免了背景太大，造成的数据集分类不均的情况。
在我的实验中，遇到的情况就是，使用了MCE，然后得到的DICE上来就是95%，然后Validation最高只能达到98.8%，predict出来和binary的对比了一下,发现效果有点差。
然后就看到⁴的一些分享，即是使用去除背景的DICE系数，并且可以加权一部分的MSE，这样既是避免了数据集分类不均衡的影像，同时又考虑了背景的分割。
好的损失函数，确实可以提高训练效率。除此之外，还可以做一些Sceneario Specific的侧重，比如在医学上，通常比较看重Recall（Sensitvity）,或者哪一种类的情况比其他更严重，就可以从损失函数中加一些超参数进行控制。
PS.
想了想自己这点水平就是做点结构工程和损失函数工程。
发现想仅仅通过损失函数对效果进行提升，还是有点力不从心。模型才是基础，损失函数更像是锦上添花，下一篇对医疗图像分割深度学习模型进行一个学习。

---

1. U-Net: Convolutional Networks for Biomedical Image Segmentation ↩︎

2. 从loss方面处理图像分割中数据极度不均衡的状况 --keras ↩︎

3. Generalised dice overlap as a deep learning loss function for highly unbalanced segmentations, ↩︎

4. KDSB 2018 第一名解决方案 ↩︎ ↩︎