138 lines
4.7 KiB
Python
138 lines
4.7 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
"""Modified from
|
|
https://github.com/JunMa11/SegLoss/blob/master/losses_pytorch/dice_loss.py#L333
|
|
(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 tversky_loss(pred,
|
|
target,
|
|
valid_mask,
|
|
alpha=0.3,
|
|
beta=0.7,
|
|
smooth=1,
|
|
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:
|
|
tversky_loss = binary_tversky_loss(
|
|
pred[:, i],
|
|
target[..., i],
|
|
valid_mask=valid_mask,
|
|
alpha=alpha,
|
|
beta=beta,
|
|
smooth=smooth)
|
|
if class_weight is not None:
|
|
tversky_loss *= class_weight[i]
|
|
total_loss += tversky_loss
|
|
return total_loss / num_classes
|
|
|
|
|
|
@weighted_loss
|
|
def binary_tversky_loss(pred,
|
|
target,
|
|
valid_mask,
|
|
alpha=0.3,
|
|
beta=0.7,
|
|
smooth=1):
|
|
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)
|
|
|
|
TP = torch.sum(torch.mul(pred, target) * valid_mask, dim=1)
|
|
FP = torch.sum(torch.mul(pred, 1 - target) * valid_mask, dim=1)
|
|
FN = torch.sum(torch.mul(1 - pred, target) * valid_mask, dim=1)
|
|
tversky = (TP + smooth) / (TP + alpha * FP + beta * FN + smooth)
|
|
|
|
return 1 - tversky
|
|
|
|
|
|
@LOSSES.register_module()
|
|
class TverskyLoss(nn.Module):
|
|
"""TverskyLoss. This loss is proposed in `Tversky loss function for image
|
|
segmentation using 3D fully convolutional deep networks.
|
|
|
|
<https://arxiv.org/abs/1706.05721>`_.
|
|
Args:
|
|
smooth (float): A float number to smooth loss, and avoid NaN error.
|
|
Default: 1.
|
|
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.
|
|
alpha(float, in [0, 1]):
|
|
The coefficient of false positives. Default: 0.3.
|
|
beta (float, in [0, 1]):
|
|
The coefficient of false negatives. Default: 0.7.
|
|
Note: alpha + beta = 1.
|
|
loss_name (str, optional): Name of the loss item. If you want this loss
|
|
item to be included into the backward graph, `loss_` must be the
|
|
prefix of the name. Defaults to 'loss_tversky'.
|
|
"""
|
|
|
|
def __init__(self,
|
|
smooth=1,
|
|
class_weight=None,
|
|
loss_weight=1.0,
|
|
ignore_index=255,
|
|
alpha=0.3,
|
|
beta=0.7,
|
|
loss_name='loss_tversky'):
|
|
super().__init__()
|
|
self.smooth = smooth
|
|
self.class_weight = get_class_weight(class_weight)
|
|
self.loss_weight = loss_weight
|
|
self.ignore_index = ignore_index
|
|
assert (alpha + beta == 1.0), 'Sum of alpha and beta but be 1.0!'
|
|
self.alpha = alpha
|
|
self.beta = beta
|
|
self._loss_name = loss_name
|
|
|
|
def forward(self, pred, target, **kwargs):
|
|
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 * tversky_loss(
|
|
pred,
|
|
one_hot_target,
|
|
valid_mask=valid_mask,
|
|
alpha=self.alpha,
|
|
beta=self.beta,
|
|
smooth=self.smooth,
|
|
class_weight=class_weight,
|
|
ignore_index=self.ignore_index)
|
|
return loss
|
|
|
|
@property
|
|
def loss_name(self):
|
|
"""Loss Name.
|
|
|
|
This function must be implemented and will return the name of this
|
|
loss function. This name will be used to combine different loss items
|
|
by simple sum operation. In addition, if you want this loss item to be
|
|
included into the backward graph, `loss_` must be the prefix of the
|
|
name.
|
|
Returns:
|
|
str: The name of this loss item.
|
|
"""
|
|
return self._loss_name
|