Spaces:
Runtime error
Runtime error
File size: 4,821 Bytes
2ae34e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
# 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
|