| |
| |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def masked_sum(loss, mask, label_weight=1, eps=1e-8, reduction=True): |
| if mask is not None: |
| loss = loss.masked_fill(mask, 0.0) |
| if reduction: |
| return loss.sum() / (((1 - mask.long()) * label_weight).sum() + eps) |
|
|
| if reduction: |
| return loss.mean() |
|
|
| return loss |
|
|
|
|
| def cross_entropy(log_prob, target, mask, focal=False, label_weight=None, reduction=True): |
| target = target.unsqueeze(-1) |
| if focal: |
| focal_coeff = log_prob.exp().gather(-1, target).squeeze(-1) |
| focal_coeff = (1.0 - focal_coeff) ** 2 |
| else: |
| focal_coeff = 1.0 |
|
|
| loss = -focal_coeff * log_prob.gather(-1, target).squeeze(-1) |
|
|
| if label_weight is not None: |
| loss = loss * label_weight |
| return masked_sum(loss, mask, label_weight=label_weight, reduction=reduction) |
| else: |
| return masked_sum(loss, mask, reduction=reduction) |
|
|
|
|
| def binary_cross_entropy(logits, target, mask, focal=False, reduction=True): |
| if focal: |
| prob = logits.sigmoid() |
| focal_coeff = target * prob + (1.0 - target) * (1.0 - prob) |
| focal_coeff = (1.0 - focal_coeff) ** 2 |
| else: |
| focal_coeff = 1.0 |
|
|
| loss = focal_coeff * F.binary_cross_entropy_with_logits(logits, target, reduction="none") |
| return masked_sum(loss, mask, reduction=reduction) |
|
|
