交叉熵损失函数好用是好用,但当数据存在严重的不平衡问题的时候,会导致结果特别糟糕,特别是在医学图像分割问题中,一般目标区域都比较小,背景区域比较大,此时用交叉熵损失函数的话得到的结果就不太好。
IOU loss介绍
IOU即是交并比,用相交的部分去除上并集,就得到IOU的值,1-IOU的值就是IOU Loss。至于IOU的数学定义去看百度百科吧,举个例子:
上面两张图求IOU,白色区域为目标区域,就是用两张图片的白色区域的交集去比上白色部分的并集,就得到了白色类别的IOU值。在实际工程中,一般黑色像素为类别0,白色为类别1。可以使用代码轻松的求出IOU值。
Pytorch代码
import numpy
import torch
import torch.nn as nn
import torch.nn.functional as F
class IoULoss(nn.Module):
def __init__(self, weight=None, size_average=True):
super(IoULoss, self).__init__()
def forward(self, inputs, targets, smooth=1):
#comment out if your model contains a sigmoid or equivalent activation layer
inputs = F.sigmoid(inputs)
#flatten label and prediction tensors
inputs = inputs.view(-1)
targets = targets.view(-1)
#intersection is equivalent to True Positive count
#union is the mutually inclusive area of all labels & predictions
intersection = (inputs * targets).sum()
total = (inputs + targets).sum()
union = total - intersection
IoU = (intersection + smooth)/(union + smooth)
return 1 - IoU