-
Notifications
You must be signed in to change notification settings - Fork 1
/
focal_loss.py
225 lines (195 loc) · 9.06 KB
/
focal_loss.py
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
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
def focal_loss(
input: torch.Tensor,
target: torch.Tensor,
alpha: float,
gamma: float = 2.0,
reduction: str = 'none',
eps: float = 1e-8) -> torch.Tensor:
r"""Criterion that computes Focal loss.
According to :cite:`lin2018focal`, the Focal loss is computed as follows:
.. math::
\text{FL}(p_t) = -\alpha_t (1 - p_t)^{\gamma} \, \text{log}(p_t)
Where:
- :math:`p_t` is the model's estimated probability for each class.
Args:
input (torch.Tensor): logits tensor with shape :math:`(N, C, *)` where C = number of classes.
target (torch.Tensor): labels tensor with shape :math:`(N, *)` where each value is :math:`0 ≤ targets[i] ≤ C−1`.
alpha (float): Weighting factor :math:`\alpha \in [0, 1]`.
gamma (float, optional): Focusing parameter :math:`\gamma >= 0`. Default 2.
reduction (str, optional): Specifies the reduction to apply to the
output: ‘none’ | ‘mean’ | ‘sum’. ‘none’: no reduction will be applied,
‘mean’: the sum of the output will be divided by the number of elements
in the output, ‘sum’: the output will be summed. Default: ‘none’.
eps (float, optional): Scalar to enforce numerical stabiliy. Default: 1e-8.
Return:
torch.Tensor: the computed loss.
Example:
>>> N = 5 # num_classes
>>> input = torch.randn(1, N, 3, 5, requires_grad=True)
>>> target = torch.empty(1, 3, 5, dtype=torch.long).random_(N)
>>> output = focal_loss(input, target, alpha=0.5, gamma=2.0, reduction='mean')
>>> output.backward()
"""
if not isinstance(input, torch.Tensor):
raise TypeError("Input type is not a torch.Tensor. Got {}"
.format(type(input)))
# if not len(input.shape) >= 2:
# raise ValueError("Invalid input shape, we expect BxCx*. Got: {}"
# .format(input.shape))
# if input.size(0) != target.size(0):
# raise ValueError('Expected input batch_size ({}) to match target batch_size ({}).'
# .format(input.size(0), target.size(0)))
# n = input.size(0)
# out_size = (n,) + input.size()[2:]
# if target.size()[1:] != input.size()[2:]:
# raise ValueError('Expected target size {}, got {}'.format(
# out_size, target.size()))
if not input.device == target.device:
raise ValueError(
"input and target must be in the same device. Got: {} and {}" .format(
input.device, target.device))
# # compute softmax over the classes axis
input_soft: torch.Tensor = F.softmax(input, dim=1) + eps
# create the labels one hot tensor
target_one_hot: torch.Tensor = one_hot(
target, num_classes=input.shape[1],
device=input.device, dtype=input.dtype)
# compute the actual focal loss
weight = torch.pow(-input_soft + 1., gamma)
focal = -alpha * weight * torch.log(input_soft)
loss_tmp = torch.sum(target_one_hot * focal, dim=1)
if reduction == 'none':
loss = loss_tmp
elif reduction == 'mean':
loss = torch.mean(loss_tmp)
elif reduction == 'sum':
loss = torch.sum(loss_tmp)
else:
raise NotImplementedError("Invalid reduction mode: {}"
.format(reduction))
return loss
class FocalLoss(nn.Module):
r"""Criterion that computes Focal loss.
According to :cite:`lin2018focal`, the Focal loss is computed as follows:
.. math::
\text{FL}(p_t) = -\alpha_t (1 - p_t)^{\gamma} \, \text{log}(p_t)
Where:
- :math:`p_t` is the model's estimated probability for each class.
Args:
alpha (float): Weighting factor :math:`\alpha \in [0, 1]`.
gamma (float, optional): Focusing parameter :math:`\gamma >= 0`. Default 2.
reduction (str, optional): Specifies the reduction to apply to the
output: ‘none’ | ‘mean’ | ‘sum’. ‘none’: no reduction will be applied,
‘mean’: the sum of the output will be divided by the number of elements
in the output, ‘sum’: the output will be summed. Default: ‘none’.
eps (float, optional): Scalar to enforce numerical stabiliy. Default: 1e-8.
Shape:
- Input: :math:`(N, C, *)` where C = number of classes.
- Target: :math:`(N, *)` where each value is
:math:`0 ≤ targets[i] ≤ C−1`.
Example:
>>> N = 5 # num_classes
>>> kwargs = {"alpha": 0.5, "gamma": 2.0, "reduction": 'mean'}
>>> criterion = FocalLoss(**kwargs)
>>> input = torch.randn(1, N, 3, 5, requires_grad=True)
>>> target = torch.empty(1, 3, 5, dtype=torch.long).random_(N)
>>> output = criterion(input, target)
>>> output.backward()
"""
def __init__(self, alpha: float, gamma: float = 2.0,
reduction: str = 'none', eps: float = 1e-8) -> None:
super(FocalLoss, self).__init__()
self.alpha: float = alpha
self.gamma: float = gamma
self.reduction: str = reduction
self.eps: float = eps
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
return focal_loss(input, target, self.alpha, self.gamma, self.reduction, self.eps)
def one_hot(labels: torch.Tensor,
num_classes: int,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
eps: float = 1e-6) -> torch.Tensor:
r"""Converts an integer label x-D tensor to a one-hot (x+1)-D tensor.
Args:
labels (torch.Tensor) : tensor with labels of shape :math:`(N, *)`,
where N is batch size. Each value is an integer
representing correct classification.
num_classes (int): number of classes in labels.
device (Optional[torch.device]): the desired device of returned tensor.
Default: if None, uses the current device for the default tensor type
(see torch.set_default_tensor_type()). device will be the CPU for CPU
tensor types and the current CUDA device for CUDA tensor types.
dtype (Optional[torch.dtype]): the desired data type of returned
tensor. Default: if None, infers data type from values.
Returns:
torch.Tensor: the labels in one hot tensor of shape :math:`(N, C, *)`,
Examples:
>>> labels = torch.LongTensor([[[0, 1], [2, 0]]])
>>> one_hot(labels, num_classes=3)
tensor([[[[1.0000e+00, 1.0000e-06],
[1.0000e-06, 1.0000e+00]],
<BLANKLINE>
[[1.0000e-06, 1.0000e+00],
[1.0000e-06, 1.0000e-06]],
<BLANKLINE>
[[1.0000e-06, 1.0000e-06],
[1.0000e+00, 1.0000e-06]]]])
"""
if not isinstance(labels, torch.Tensor):
raise TypeError("Input labels type is not a torch.Tensor. Got {}"
.format(type(labels)))
if not labels.dtype == torch.int64:
raise ValueError(
"labels must be of the same dtype torch.int64. Got: {}" .format(
labels.dtype))
if num_classes < 1:
raise ValueError("The number of classes must be bigger than one."
" Got: {}".format(num_classes))
shape = labels.shape
one_hot = torch.zeros(
(shape[0], num_classes) + shape[1:], device=device, dtype=dtype
)
return one_hot.scatter_(1, labels.unsqueeze(1), 1.0) + eps
# import torch
# import torch.nn as nn
# class FocalLoss(nn.Module):
# def __init__(self, alpha=1, gamma=2, reduction: str = 'mean'):
# super().__init__()
# if reduction not in ['mean', 'none', 'sum']:
# raise NotImplementedError('Reduction {} not implemented.'.format(reduction))
# self.reduction = reduction
# self.alpha = alpha
# self.gamma = gamma
# def forward(self, x, target):
# p_t = torch.where(target == 1, x, 1-x)
# fl = - 1 * (1 - p_t) ** self.gamma * torch.log(p_t)
# fl = torch.where(target == 1, fl * self.alpha, fl)
# return self._reduce(fl)
# def _reduce(self, x):
# if self.reduction == 'mean':
# return x.mean()
# elif self.reduction == 'sum':
# return x.sum()
# else:
# return x
# import torch
# import torch.nn as nn
# import torch.functional as F
# class FocalLoss(nn.modules.loss._WeightedLoss):
# r'''
# FL(p_t) = -(1-p_t)^\gamma \log(p_t)
# '''
# def __init__(self, weight=None, gamma=2, reduction='mean'):
# super(FocalLoss, self).__init__(weight,reduction=reduction)
# self.gamma = gamma
# self.weight = weight #weight parameter will act as the alpha parameter to balance class weights
# def forward(self, data, target):
# ce_loss = F.cross_entropy(data, target, reduction=self.reduction,weight=self.weight)
# pt = torch.exp(-ce_loss)
# focal_loss = ((1 - pt) ** self.gamma * ce_loss).mean()
# return focal_loss