Spaces:
Runtime error
Runtime error
File size: 2,528 Bytes
8075387 |
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 |
import torch
from torch.nn import functional as F
# TODO: merge these two function
def heatmap_focal_loss(
inputs,
targets,
pos_inds,
labels,
alpha: float = -1,
beta: float = 4,
gamma: float = 2,
reduction: str = 'sum',
sigmoid_clamp: float = 1e-4,
ignore_high_fp: float = -1.,
):
"""
Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
Args:
inputs: (sum_l N*Hl*Wl, C)
targets: (sum_l N*Hl*Wl, C)
pos_inds: N
labels: N
Returns:
Loss tensor with the reduction option applied.
"""
pred = torch.clamp(inputs.sigmoid_(), min=sigmoid_clamp, max=1-sigmoid_clamp)
neg_weights = torch.pow(1 - targets, beta)
pos_pred_pix = pred[pos_inds] # N x C
pos_pred = pos_pred_pix.gather(1, labels.unsqueeze(1))
pos_loss = torch.log(pos_pred) * torch.pow(1 - pos_pred, gamma)
neg_loss = torch.log(1 - pred) * torch.pow(pred, gamma) * neg_weights
if ignore_high_fp > 0:
not_high_fp = (pred < ignore_high_fp).float()
neg_loss = not_high_fp * neg_loss
if reduction == "sum":
pos_loss = pos_loss.sum()
neg_loss = neg_loss.sum()
if alpha >= 0:
pos_loss = alpha * pos_loss
neg_loss = (1 - alpha) * neg_loss
return - pos_loss, - neg_loss
heatmap_focal_loss_jit = torch.jit.script(heatmap_focal_loss)
# heatmap_focal_loss_jit = heatmap_focal_loss
def binary_heatmap_focal_loss(
inputs,
targets,
pos_inds,
alpha: float = -1,
beta: float = 4,
gamma: float = 2,
sigmoid_clamp: float = 1e-4,
ignore_high_fp: float = -1.,
):
"""
Args:
inputs: (sum_l N*Hl*Wl,)
targets: (sum_l N*Hl*Wl,)
pos_inds: N
Returns:
Loss tensor with the reduction option applied.
"""
pred = torch.clamp(inputs.sigmoid_(), min=sigmoid_clamp, max=1-sigmoid_clamp)
neg_weights = torch.pow(1 - targets, beta)
pos_pred = pred[pos_inds] # N
pos_loss = torch.log(pos_pred) * torch.pow(1 - pos_pred, gamma)
neg_loss = torch.log(1 - pred) * torch.pow(pred, gamma) * neg_weights
if ignore_high_fp > 0:
not_high_fp = (pred < ignore_high_fp).float()
neg_loss = not_high_fp * neg_loss
pos_loss = - pos_loss.sum()
neg_loss = - neg_loss.sum()
if alpha >= 0:
pos_loss = alpha * pos_loss
neg_loss = (1 - alpha) * neg_loss
return pos_loss, neg_loss
binary_heatmap_focal_loss_jit = torch.jit.script(binary_heatmap_focal_loss) |