Spaces:
Runtime error
Runtime error
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 | |