|
"""Modified from https://github.com/LikeLy-Journey/SegmenTron/blob/master/ |
|
segmentron/solver/loss.py (Apache-2.0 License)""" |
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
from ..builder import LOSSES |
|
from .utils import get_class_weight, weighted_loss |
|
|
|
|
|
@weighted_loss |
|
def dice_loss(pred, |
|
target, |
|
valid_mask, |
|
smooth=1, |
|
exponent=2, |
|
class_weight=None, |
|
ignore_index=255): |
|
assert pred.shape[0] == target.shape[0] |
|
total_loss = 0 |
|
num_classes = pred.shape[1] |
|
for i in range(num_classes): |
|
if i != ignore_index: |
|
dice_loss = binary_dice_loss( |
|
pred[:, i], |
|
target[..., i], |
|
valid_mask=valid_mask, |
|
smooth=smooth, |
|
exponent=exponent) |
|
if class_weight is not None: |
|
dice_loss *= class_weight[i] |
|
total_loss += dice_loss |
|
return total_loss / num_classes |
|
|
|
|
|
@weighted_loss |
|
def binary_dice_loss(pred, target, valid_mask, smooth=1, exponent=2, **kwards): |
|
assert pred.shape[0] == target.shape[0] |
|
pred = pred.reshape(pred.shape[0], -1) |
|
target = target.reshape(target.shape[0], -1) |
|
valid_mask = valid_mask.reshape(valid_mask.shape[0], -1) |
|
|
|
num = torch.sum(torch.mul(pred, target) * valid_mask, dim=1) * 2 + smooth |
|
den = torch.sum(pred.pow(exponent) + target.pow(exponent), dim=1) + smooth |
|
|
|
return 1 - num / den |
|
|
|
|
|
@LOSSES.register_module() |
|
class DiceLoss(nn.Module): |
|
"""DiceLoss. |
|
|
|
This loss is proposed in `V-Net: Fully Convolutional Neural Networks for |
|
Volumetric Medical Image Segmentation <https://arxiv.org/abs/1606.04797>`_. |
|
|
|
Args: |
|
loss_type (str, optional): Binary or multi-class loss. |
|
Default: 'multi_class'. Options are "binary" and "multi_class". |
|
smooth (float): A float number to smooth loss, and avoid NaN error. |
|
Default: 1 |
|
exponent (float): An float number to calculate denominator |
|
value: \\sum{x^exponent} + \\sum{y^exponent}. Default: 2. |
|
reduction (str, optional): The method used to reduce the loss. Options |
|
are "none", "mean" and "sum". This parameter only works when |
|
per_image is True. Default: 'mean'. |
|
class_weight (list[float] | str, optional): Weight of each class. If in |
|
str format, read them from a file. Defaults to None. |
|
loss_weight (float, optional): Weight of the loss. Default to 1.0. |
|
ignore_index (int | None): The label index to be ignored. Default: 255. |
|
""" |
|
|
|
def __init__(self, |
|
smooth=1, |
|
exponent=2, |
|
reduction='mean', |
|
class_weight=None, |
|
loss_weight=1.0, |
|
ignore_index=255, |
|
**kwards): |
|
super(DiceLoss, self).__init__() |
|
self.smooth = smooth |
|
self.exponent = exponent |
|
self.reduction = reduction |
|
self.class_weight = get_class_weight(class_weight) |
|
self.loss_weight = loss_weight |
|
self.ignore_index = ignore_index |
|
|
|
def forward(self, |
|
pred, |
|
target, |
|
avg_factor=None, |
|
reduction_override=None, |
|
**kwards): |
|
assert reduction_override in (None, 'none', 'mean', 'sum') |
|
reduction = ( |
|
reduction_override if reduction_override else self.reduction) |
|
if self.class_weight is not None: |
|
class_weight = pred.new_tensor(self.class_weight) |
|
else: |
|
class_weight = None |
|
|
|
pred = F.softmax(pred, dim=1) |
|
num_classes = pred.shape[1] |
|
one_hot_target = F.one_hot( |
|
torch.clamp(target.long(), 0, num_classes - 1), |
|
num_classes=num_classes) |
|
valid_mask = (target != self.ignore_index).long() |
|
|
|
loss = self.loss_weight * dice_loss( |
|
pred, |
|
one_hot_target, |
|
valid_mask=valid_mask, |
|
reduction=reduction, |
|
avg_factor=avg_factor, |
|
smooth=self.smooth, |
|
exponent=self.exponent, |
|
class_weight=class_weight, |
|
ignore_index=self.ignore_index) |
|
return loss |
|
|
