pytorch 多类分割损失 （Generalized Dice Loss）

引自该文章python 常用函数和自定义函数整理 ‘2. Pytorch相关处理’Generalized Dice Loss相关代码，如有错误，烦请指正。

# 多类分割dice损失
def generalized_dice_loss(pred, target):
    """compute the weighted dice_loss
    Args:
        pred (tensor): prediction after softmax, shape(bath_size, channels, height, width)
        target (tensor): gt, shape(bath_size, channels, height, width)
    Returns:
        gldice_loss: loss value
    """    
    wei = torch.sum(target, axis=[0,2,3]) # (n_class,)
    wei = 1/(wei**2+epsilon)
    intersection = torch.sum(wei*torch.sum(pred * target, axis=[0,2,3]))
    union = torch.sum(wei*torch.sum(pred + target, axis=[0,2,3]))
    gldice_loss = 1 - (2. * intersection) / (union + epsilon)
    return gldice_loss

# 多类交叉熵损失
import torch.nn.functional as F
def BCE_loss(pred, target):
    """compute the bce_loss 
    Args:
        pred (tensor): prediction after softmax, shape(N, C, H, W)， channels = class numbers
        target (tensor): gt, shape(N, H, W), each value should be between [0,C)
    Returns:
        bce_loss : loss value
    """    
    target= target.long().to(self.device)
    bce_loss = F.cross_entropy(pred, target)
    
    return bce_loss