File size: 5,732 Bytes
26225c5 | 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | from typing import Optional, Sequence
import torch
from torch import Tensor
from torch import nn
from torch.nn import functional as F
__all__ = ['WeightedFocalLoss']
class WeightedFocalLoss(nn.Module):
""" Focal Loss, as described in https://arxiv.org/abs/1708.02002.
It is essentially an enhancement to cross entropy loss and is
useful for classification tasks when there is a large class imbalance.
x is expected to contain raw, unnormalized scores for each class.
y is expected to contain class labels.
Shape:
- x: (batch_size, C) or (batch_size, C, d1, d2, ..., dK), K > 0.
- y: (batch_size,) or (batch_size, d1, d2, ..., dK), K > 0.
Credit: https://github.com/AdeelH/pytorch-multi-class-focal-loss
Note:
Modified `alpha` to `weight` to respect the loss template
expected by `loss_with_target_histogram`
"""
def __init__(
self,
weight: Optional[Tensor] = None,
gamma: float = 0.,
reduction: str = 'mean',
ignore_index: int = -100):
"""Constructor.
Args:
weight (Tensor): Weights for each class. Defaults to None.
gamma (float): A constant, as described in the paper.
Defaults to 0.
reduction (str): 'mean' or 'none'.
Defaults to 'mean'.
ignore_index (int): class label to ignore.
Defaults to -100.
"""
if reduction not in ('mean', 'none'):
raise ValueError(
'Reduction must be one of: "mean", "none".')
super().__init__()
self.weight = weight
self.gamma = gamma
self.ignore_index = ignore_index
self.reduction = reduction
self.nll_loss = nn.NLLLoss(
weight=weight, reduction='none', ignore_index=ignore_index)
def __repr__(self):
arg_keys = ['weight', 'gamma', 'ignore_index', 'reduction']
arg_vals = [self.__dict__[k] for k in arg_keys]
arg_strs = [f'{k}={v!r}' for k, v in zip(arg_keys, arg_vals)]
arg_str = ', '.join(arg_strs)
return f'{type(self).__name__}({arg_str})'
def forward(self, x: Tensor, y: Tensor, w: Tensor) -> Tensor:
"""
:param x: (N, C, ...) Tensor
Logits
:param y: (N, ...) Tensor
Target labels
:param w: (N, ...) Tensor
Per-item weights, can be None
"""
# Convert 1D x to multiclass x. This is an artificial step for
# binary classification (eg affinity loss), where only 1 score
# is provided. In this case, we assume that x<0 accounts for y=0
# and x>0 accounts for y=1 (ie prepared for sigmoid). Here, we
# convert these precitions to 2D for downstream softmax
if x.dim() == 1:
x_binary = torch.zeros(x.shape[0], 2, dtype=x.dtype, device=x.device)
x_binary[x < 0, 0] = -x[x < 0]
x_binary[x > 0, 1] = x[x > 0]
x = x_binary
# Convert y to long. The NLL loss does not support non-integer
# target labels
y = y.long()
# Convert per-item weights to [0, 1] weights
if w is None:
w = torch.ones_like(y).float()
w = w / w.sum()
if x.ndim > 2:
# (N, C, d1, d2, ..., dK) --> (N * d1 * ... * dK, C)
c = x.shape[1]
x = x.permute(0, *range(2, x.ndim), 1).reshape(-1, c)
# (N, d1, d2, ..., dK) --> (N * d1 * ... * dK,)
y = y.view(-1)
# (N, d1, d2, ..., dK) --> (N * d1 * ... * dK,)
w = w.view(-1)
unignored_mask = y != self.ignore_index
y = y[unignored_mask]
if len(y) == 0:
return torch.tensor(0., device=x.device)
x = x[unignored_mask]
w = w[unignored_mask]
# compute weighted cross entropy term: -weight * log(pt)
# (weight is already part of self.nll_loss)
log_p = F.log_softmax(x, dim=-1)
ce = self.nll_loss(log_p, y)
# get true class column from each row
log_pt = log_p.gather(dim=1, index=y.view(-1, 1)).squeeze()
# compute focal term: (1 - pt)^gamma
pt = log_pt.exp()
focal_term = (1 - pt) ** self.gamma
# the full loss: -weight * ((1 - pt)^gamma) * log(pt)
loss = focal_term * ce
# Apply the per-item weighting
loss = loss * w
if self.reduction == 'none':
return loss
return loss.sum()
def weighted_focal_loss(
weight: Optional[Sequence] = None,
gamma: float = 0.,
reduction: str = 'mean',
ignore_index: int = -100,
device='cpu',
dtype=torch.float32) -> WeightedFocalLoss:
"""Factory function for WeightedFocalLoss.
Args:
weight (Sequence): Weights for each class. Will be converted
to a Tensor if not None. Defaults to None.
gamma (float): A constant, as described in the paper.
Defaults to 0.
reduction (str): 'mean' or 'none'.
Defaults to 'mean'.
ignore_index (int): class label to ignore.
Defaults to -100.
device (str): Device to move weight to. Defaults to 'cpu'.
dtype (torch.dtype): dtype to cast weight to.
Defaults to torch.float32.
Returns:
A WeightedFocalLoss object
"""
if weight is not None:
if not isinstance(weight, Tensor):
weight = torch.tensor(weight)
weight = weight.to(device=device, dtype=dtype)
fl = WeightedFocalLoss(
weight=weight,
gamma=gamma,
reduction=reduction,
ignore_index=ignore_index)
return fl
|