You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
98 lines
3.2 KiB
98 lines
3.2 KiB
4 years ago
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
|
||
|
|
|
||
|
|
class AsymmetricLossMultiLabel(nn.Module):
|
||
|
|
def __init__(self, gamma_neg=4, gamma_pos=1, clip=0.05, eps=1e-8, disable_torch_grad_focal_loss=False):
|
||
|
|
super(AsymmetricLossMultiLabel, self).__init__()
|
||
|
|
|
||
|
|
self.gamma_neg = gamma_neg
|
||
|
|
self.gamma_pos = gamma_pos
|
||
|
|
self.clip = clip
|
||
|
|
self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
|
||
|
|
self.eps = eps
|
||
|
|
|
||
|
|
def forward(self, x, y):
|
||
|
|
""""
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
x: input logits
|
||
|
|
y: targets (multi-label binarized vector)
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Calculating Probabilities
|
||
|
|
x_sigmoid = torch.sigmoid(x)
|
||
|
|
xs_pos = x_sigmoid
|
||
|
|
xs_neg = 1 - x_sigmoid
|
||
|
|
|
||
|
|
# Asymmetric Clipping
|
||
|
|
if self.clip is not None and self.clip > 0:
|
||
|
|
xs_neg = (xs_neg + self.clip).clamp(max=1)
|
||
|
|
|
||
|
|
# Basic CE calculation
|
||
|
|
los_pos = y * torch.log(xs_pos.clamp(min=self.eps))
|
||
|
|
los_neg = (1 - y) * torch.log(xs_neg.clamp(min=self.eps))
|
||
|
|
loss = los_pos + los_neg
|
||
|
|
|
||
|
|
# Asymmetric Focusing
|
||
|
|
if self.gamma_neg > 0 or self.gamma_pos > 0:
|
||
|
|
if self.disable_torch_grad_focal_loss:
|
||
|
|
torch._C.set_grad_enabled(False)
|
||
|
|
pt0 = xs_pos * y
|
||
|
|
pt1 = xs_neg * (1 - y) # pt = p if t > 0 else 1-p
|
||
|
|
pt = pt0 + pt1
|
||
|
|
one_sided_gamma = self.gamma_pos * y + self.gamma_neg * (1 - y)
|
||
|
|
one_sided_w = torch.pow(1 - pt, one_sided_gamma)
|
||
|
|
if self.disable_torch_grad_focal_loss:
|
||
|
|
torch._C.set_grad_enabled(True)
|
||
|
|
loss *= one_sided_w
|
||
|
|
|
||
|
|
return -loss.sum()
|
||
|
|
|
||
|
|
|
||
|
|
class AsymmetricLossSingleLabel(nn.Module):
|
||
|
|
def __init__(self, gamma_pos=1, gamma_neg=4, eps: float = 0.1, reduction='mean'):
|
||
|
|
super(AsymmetricLossSingleLabel, self).__init__()
|
||
|
|
|
||
|
|
self.eps = eps
|
||
|
|
self.logsoftmax = nn.LogSoftmax(dim=-1)
|
||
|
|
self.targets_classes = [] # prevent gpu repeated memory allocation
|
||
|
|
self.gamma_pos = gamma_pos
|
||
|
|
self.gamma_neg = gamma_neg
|
||
|
|
self.reduction = reduction
|
||
|
|
|
||
|
|
def forward(self, inputs, target, reduction=None):
|
||
|
|
""""
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
x: input logits
|
||
|
|
y: targets (1-hot vector)
|
||
|
|
"""
|
||
|
|
|
||
|
|
num_classes = inputs.size()[-1]
|
||
|
|
log_preds = self.logsoftmax(inputs)
|
||
|
|
self.targets_classes = torch.zeros_like(inputs).scatter_(1, target.long().unsqueeze(1), 1)
|
||
|
|
|
||
|
|
# ASL weights
|
||
|
|
targets = self.targets_classes
|
||
|
|
anti_targets = 1 - targets
|
||
|
|
xs_pos = torch.exp(log_preds)
|
||
|
|
xs_neg = 1 - xs_pos
|
||
|
|
xs_pos = xs_pos * targets
|
||
|
|
xs_neg = xs_neg * anti_targets
|
||
|
|
asymmetric_w = torch.pow(1 - xs_pos - xs_neg,
|
||
|
|
self.gamma_pos * targets + self.gamma_neg * anti_targets)
|
||
|
|
log_preds = log_preds * asymmetric_w
|
||
|
|
|
||
|
|
if self.eps > 0: # label smoothing
|
||
|
|
self.targets_classes.mul_(1 - self.eps).add_(self.eps / num_classes)
|
||
|
|
|
||
|
|
# loss calculation
|
||
|
|
loss = - self.targets_classes.mul(log_preds)
|
||
|
|
|
||
|
|
loss = loss.sum(dim=-1)
|
||
|
|
if self.reduction == 'mean':
|
||
|
|
loss = loss.mean()
|
||
|
|
|
||
|
|
return loss
|