mirror of https://github.com/alibaba/EasyCV.git
388 lines
16 KiB
Python
388 lines
16 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from mmcv.ops import sigmoid_focal_loss as _sigmoid_focal_loss
|
|
|
|
from easycv.models.builder import LOSSES
|
|
from easycv.models.loss.utils import weight_reduce_loss
|
|
|
|
|
|
def reduce_loss(loss, reduction):
|
|
"""Reduce loss as specified.
|
|
Args:
|
|
loss (Tensor): Elementwise loss tensor.
|
|
reduction (str): Options are "none", "mean" and "sum".
|
|
Return:
|
|
Tensor: Reduced loss tensor.
|
|
"""
|
|
reduction_enum = F._Reduction.get_enum(reduction)
|
|
# none: 0, elementwise_mean:1, sum: 2
|
|
if reduction_enum == 0:
|
|
return loss
|
|
elif reduction_enum == 1:
|
|
return loss.mean()
|
|
elif reduction_enum == 2:
|
|
return loss.sum()
|
|
else:
|
|
raise ValueError('reduction_enum should be 0,1,2')
|
|
|
|
|
|
def varifocal_loss(pred,
|
|
target,
|
|
weight=None,
|
|
alpha=0.75,
|
|
gamma=2.0,
|
|
iou_weighted=True,
|
|
reduction='mean',
|
|
avg_factor=None):
|
|
"""`Varifocal Loss <https://arxiv.org/abs/2008.13367>`_
|
|
Args:
|
|
pred (torch.Tensor): The prediction with shape (N, C), C is the
|
|
number of classes
|
|
target (torch.Tensor): The learning target of the iou-aware
|
|
classification score with shape (N, C), C is the number of classes.
|
|
weight (torch.Tensor, optional): The weight of loss for each
|
|
prediction. Defaults to None.
|
|
alpha (float, optional): A balance factor for the negative part of
|
|
Varifocal Loss, which is different from the alpha of Focal Loss.
|
|
Defaults to 0.75.
|
|
gamma (float, optional): The gamma for calculating the modulating
|
|
factor. Defaults to 2.0.
|
|
iou_weighted (bool, optional): Whether to weight the loss of the
|
|
positive example with the iou target. Defaults to True.
|
|
reduction (str, optional): The method used to reduce the loss into
|
|
a scalar. Defaults to 'mean'. Options are "none", "mean" and
|
|
"sum".
|
|
avg_factor (int, optional): Average factor that is used to average
|
|
the loss. Defaults to None.
|
|
"""
|
|
# pred and target should be of the same size
|
|
assert pred.size() == target.size()
|
|
pred_sigmoid = pred.sigmoid()
|
|
target = target.type_as(pred)
|
|
if iou_weighted:
|
|
focal_weight = target * (target > 0.0).float() + \
|
|
alpha * (pred_sigmoid - target).abs().pow(gamma) * \
|
|
(target <= 0.0).float()
|
|
else:
|
|
focal_weight = (target > 0.0).float() + \
|
|
alpha * (pred_sigmoid - target).abs().pow(gamma) * \
|
|
(target <= 0.0).float()
|
|
loss = F.binary_cross_entropy_with_logits(
|
|
pred, target, reduction='none') * focal_weight
|
|
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
|
|
return loss
|
|
|
|
|
|
@LOSSES.register_module
|
|
class VarifocalLoss(nn.Module):
|
|
|
|
def __init__(self,
|
|
use_sigmoid=True,
|
|
alpha=0.75,
|
|
gamma=2.0,
|
|
iou_weighted=True,
|
|
reduction='mean',
|
|
loss_weight=1.0):
|
|
"""`Varifocal Loss <https://arxiv.org/abs/2008.13367>`_
|
|
Args:
|
|
use_sigmoid (bool, optional): Whether the prediction is
|
|
used for sigmoid or softmax. Defaults to True.
|
|
alpha (float, optional): A balance factor for the negative part of
|
|
Varifocal Loss, which is different from the alpha of Focal
|
|
Loss. Defaults to 0.75.
|
|
gamma (float, optional): The gamma for calculating the modulating
|
|
factor. Defaults to 2.0.
|
|
iou_weighted (bool, optional): Whether to weight the loss of the
|
|
positive examples with the iou target. Defaults to True.
|
|
reduction (str, optional): The method used to reduce the loss into
|
|
a scalar. Defaults to 'mean'. Options are "none", "mean" and
|
|
"sum".
|
|
loss_weight (float, optional): Weight of loss. Defaults to 1.0.
|
|
"""
|
|
super(VarifocalLoss, self).__init__()
|
|
assert use_sigmoid is True, \
|
|
'Only sigmoid varifocal loss supported now.'
|
|
assert alpha >= 0.0
|
|
self.use_sigmoid = use_sigmoid
|
|
self.alpha = alpha
|
|
self.gamma = gamma
|
|
self.iou_weighted = iou_weighted
|
|
self.reduction = reduction
|
|
self.loss_weight = loss_weight
|
|
|
|
def forward(self,
|
|
pred,
|
|
target,
|
|
weight=None,
|
|
avg_factor=None,
|
|
reduction_override=None):
|
|
"""Forward function.
|
|
Args:
|
|
pred (torch.Tensor): The prediction.
|
|
target (torch.Tensor): The learning target of the prediction.
|
|
weight (torch.Tensor, optional): The weight of loss for each
|
|
prediction. Defaults to None.
|
|
avg_factor (int, optional): Average factor that is used to average
|
|
the loss. Defaults to None.
|
|
reduction_override (str, optional): The reduction method used to
|
|
override the original reduction method of the loss.
|
|
Options are "none", "mean" and "sum".
|
|
Returns:
|
|
torch.Tensor: The calculated loss
|
|
"""
|
|
assert reduction_override in (None, 'none', 'mean', 'sum')
|
|
reduction = (
|
|
reduction_override if reduction_override else self.reduction)
|
|
if self.use_sigmoid:
|
|
loss_cls = self.loss_weight * varifocal_loss(
|
|
pred,
|
|
target,
|
|
weight,
|
|
alpha=self.alpha,
|
|
gamma=self.gamma,
|
|
iou_weighted=self.iou_weighted,
|
|
reduction=reduction,
|
|
avg_factor=avg_factor)
|
|
else:
|
|
raise NotImplementedError
|
|
return loss_cls
|
|
|
|
|
|
def py_sigmoid_focal_loss(inputs,
|
|
targets,
|
|
weight=None,
|
|
gamma=2.0,
|
|
alpha=0.25,
|
|
reduction='mean',
|
|
avg_factor=None):
|
|
"""PyTorch version of `Focal Loss <https://arxiv.org/abs/1708.02002>`_.
|
|
|
|
Args:
|
|
inputs (torch.Tensor): The prediction with shape (N, C), C is the
|
|
number of classes
|
|
targets (torch.Tensor): The learning label of the prediction.
|
|
weight (torch.Tensor, optional): Sample-wise loss weight.
|
|
gamma (float, optional): The gamma for calculating the modulating
|
|
factor. Defaults to 2.0.
|
|
alpha (float, optional): A balanced form for Focal Loss.
|
|
Defaults to 0.25.
|
|
reduction (str, optional): The method used to reduce the loss into
|
|
a scalar. Defaults to 'mean'.
|
|
avg_factor (int, optional): Average factor that is used to average
|
|
the loss. Defaults to None.
|
|
"""
|
|
prob = inputs.sigmoid()
|
|
ce_loss = F.binary_cross_entropy_with_logits(
|
|
inputs, targets, reduction='none')
|
|
p_t = prob * targets + (1 - prob) * (1 - targets)
|
|
loss = ce_loss * ((1 - p_t)**gamma)
|
|
|
|
if alpha >= 0:
|
|
alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
|
|
loss = alpha_t * loss
|
|
if weight is not None:
|
|
if weight.shape != loss.shape:
|
|
if weight.size(0) == loss.size(0):
|
|
# For most cases, weight is of shape (num_priors, ),
|
|
# which means it does not have the second axis num_class
|
|
weight = weight.view(-1, 1)
|
|
else:
|
|
# Sometimes, weight per anchor per class is also needed. e.g.
|
|
# in FSAF. But it may be flattened of shape
|
|
# (num_priors x num_class, ), while loss is still of shape
|
|
# (num_priors, num_class).
|
|
assert weight.numel() == loss.numel()
|
|
weight = weight.view(loss.size(0), -1)
|
|
assert weight.ndim == loss.ndim
|
|
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
|
|
return loss
|
|
|
|
|
|
def py_focal_loss_with_prob(pred,
|
|
target,
|
|
weight=None,
|
|
gamma=2.0,
|
|
alpha=0.25,
|
|
reduction='mean',
|
|
avg_factor=None):
|
|
"""PyTorch version of `Focal Loss <https://arxiv.org/abs/1708.02002>`_.
|
|
Different from `py_sigmoid_focal_loss`, this function accepts probability
|
|
as input.
|
|
|
|
Args:
|
|
pred (torch.Tensor): The prediction probability with shape (N, C),
|
|
C is the number of classes.
|
|
target (torch.Tensor): The learning label of the prediction.
|
|
weight (torch.Tensor, optional): Sample-wise loss weight.
|
|
gamma (float, optional): The gamma for calculating the modulating
|
|
factor. Defaults to 2.0.
|
|
alpha (float, optional): A balanced form for Focal Loss.
|
|
Defaults to 0.25.
|
|
reduction (str, optional): The method used to reduce the loss into
|
|
a scalar. Defaults to 'mean'.
|
|
avg_factor (int, optional): Average factor that is used to average
|
|
the loss. Defaults to None.
|
|
"""
|
|
num_classes = pred.size(1)
|
|
target = F.one_hot(target, num_classes=num_classes + 1)
|
|
target = target[:, :num_classes]
|
|
|
|
target = target.type_as(pred)
|
|
pt = (1 - pred) * target + pred * (1 - target)
|
|
focal_weight = (alpha * target + (1 - alpha) *
|
|
(1 - target)) * pt.pow(gamma)
|
|
loss = F.binary_cross_entropy(
|
|
pred, target, reduction='none') * focal_weight
|
|
if weight is not None:
|
|
if weight.shape != loss.shape:
|
|
if weight.size(0) == loss.size(0):
|
|
# For most cases, weight is of shape (num_priors, ),
|
|
# which means it does not have the second axis num_class
|
|
weight = weight.view(-1, 1)
|
|
else:
|
|
# Sometimes, weight per anchor per class is also needed. e.g.
|
|
# in FSAF. But it may be flattened of shape
|
|
# (num_priors x num_class, ), while loss is still of shape
|
|
# (num_priors, num_class).
|
|
assert weight.numel() == loss.numel()
|
|
weight = weight.view(loss.size(0), -1)
|
|
assert weight.ndim == loss.ndim
|
|
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
|
|
return loss
|
|
|
|
|
|
def sigmoid_focal_loss(pred,
|
|
target,
|
|
weight=None,
|
|
gamma=2.0,
|
|
alpha=0.25,
|
|
reduction='mean',
|
|
avg_factor=None):
|
|
r"""A warpper of cuda version `Focal Loss
|
|
<https://arxiv.org/abs/1708.02002>`_.
|
|
|
|
Args:
|
|
pred (torch.Tensor): The prediction with shape (N, C), C is the number
|
|
of classes.
|
|
target (torch.Tensor): The learning label of the prediction.
|
|
weight (torch.Tensor, optional): Sample-wise loss weight.
|
|
gamma (float, optional): The gamma for calculating the modulating
|
|
factor. Defaults to 2.0.
|
|
alpha (float, optional): A balanced form for Focal Loss.
|
|
Defaults to 0.25.
|
|
reduction (str, optional): The method used to reduce the loss into
|
|
a scalar. Defaults to 'mean'. Options are "none", "mean" and "sum".
|
|
avg_factor (int, optional): Average factor that is used to average
|
|
the loss. Defaults to None.
|
|
"""
|
|
# Function.apply does not accept keyword arguments, so the decorator
|
|
# "weighted_loss" is not applicable
|
|
loss = _sigmoid_focal_loss(pred.contiguous(), target.contiguous(), gamma,
|
|
alpha, None, 'none')
|
|
if weight is not None:
|
|
if weight.shape != loss.shape:
|
|
if weight.size(0) == loss.size(0):
|
|
# For most cases, weight is of shape (num_priors, ),
|
|
# which means it does not have the second axis num_class
|
|
weight = weight.view(-1, 1)
|
|
else:
|
|
# Sometimes, weight per anchor per class is also needed. e.g.
|
|
# in FSAF. But it may be flattened of shape
|
|
# (num_priors x num_class, ), while loss is still of shape
|
|
# (num_priors, num_class).
|
|
assert weight.numel() == loss.numel()
|
|
weight = weight.view(loss.size(0), -1)
|
|
assert weight.ndim == loss.ndim
|
|
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
|
|
return loss
|
|
|
|
|
|
@LOSSES.register_module()
|
|
class FocalLoss(nn.Module):
|
|
|
|
def __init__(self,
|
|
use_sigmoid=True,
|
|
gamma=2.0,
|
|
alpha=0.25,
|
|
reduction='mean',
|
|
loss_weight=1.0,
|
|
activated=False):
|
|
"""`Focal Loss <https://arxiv.org/abs/1708.02002>`_
|
|
|
|
Args:
|
|
use_sigmoid (bool, optional): Whether to the prediction is
|
|
used for sigmoid or softmax. Defaults to True.
|
|
gamma (float, optional): The gamma for calculating the modulating
|
|
factor. Defaults to 2.0.
|
|
alpha (float, optional): A balanced form for Focal Loss.
|
|
Defaults to 0.25.
|
|
reduction (str, optional): The method used to reduce the loss into
|
|
a scalar. Defaults to 'mean'. Options are "none", "mean" and
|
|
"sum".
|
|
loss_weight (float, optional): Weight of loss. Defaults to 1.0.
|
|
activated (bool, optional): Whether the input is activated.
|
|
If True, it means the input has been activated and can be
|
|
treated as probabilities. Else, it should be treated as logits.
|
|
Defaults to False.
|
|
"""
|
|
super(FocalLoss, self).__init__()
|
|
assert use_sigmoid is True, 'Only sigmoid focal loss supported now.'
|
|
self.use_sigmoid = use_sigmoid
|
|
self.gamma = gamma
|
|
self.alpha = alpha
|
|
self.reduction = reduction
|
|
self.loss_weight = loss_weight
|
|
self.activated = activated
|
|
|
|
def forward(self,
|
|
pred,
|
|
target,
|
|
weight=None,
|
|
avg_factor=None,
|
|
reduction_override=None):
|
|
"""Forward function.
|
|
|
|
Args:
|
|
pred (torch.Tensor): The prediction.
|
|
target (torch.Tensor): The learning label of the prediction.
|
|
weight (torch.Tensor, optional): The weight of loss for each
|
|
prediction. Defaults to None.
|
|
avg_factor (int, optional): Average factor that is used to average
|
|
the loss. Defaults to None.
|
|
reduction_override (str, optional): The reduction method used to
|
|
override the original reduction method of the loss.
|
|
Options are "none", "mean" and "sum".
|
|
|
|
Returns:
|
|
torch.Tensor: The calculated loss
|
|
"""
|
|
assert reduction_override in (None, 'none', 'mean', 'sum')
|
|
reduction = (
|
|
reduction_override if reduction_override else self.reduction)
|
|
if self.use_sigmoid:
|
|
if self.activated:
|
|
calculate_loss_func = py_focal_loss_with_prob
|
|
else:
|
|
if torch.cuda.is_available() and pred.is_cuda:
|
|
calculate_loss_func = sigmoid_focal_loss
|
|
else:
|
|
num_classes = pred.size(1)
|
|
target = F.one_hot(target, num_classes=num_classes + 1)
|
|
target = target[:, :num_classes]
|
|
calculate_loss_func = py_sigmoid_focal_loss
|
|
|
|
loss_cls = self.loss_weight * calculate_loss_func(
|
|
pred,
|
|
target,
|
|
weight,
|
|
gamma=self.gamma,
|
|
alpha=self.alpha,
|
|
reduction=reduction,
|
|
avg_factor=avg_factor)
|
|
|
|
else:
|
|
raise NotImplementedError
|
|
return loss_cls
|