引自该文章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