smooth_l1_loss.py 1.26 KB
Newer Older
Kai Chen's avatar
Kai Chen committed
1
import torch
Jiangmiao Pang's avatar
Jiangmiao Pang committed
2
3
import torch.nn as nn

Kai Chen's avatar
Kai Chen committed
4
from .utils import weighted_loss
Jiangmiao Pang's avatar
Jiangmiao Pang committed
5
6
7
from ..registry import LOSSES


Kai Chen's avatar
Kai Chen committed
8
9
10
11
12
13
14
15
16
17
@weighted_loss
def smooth_l1_loss(pred, target, beta=1.0):
    assert beta > 0
    assert pred.size() == target.size() and target.numel() > 0
    diff = torch.abs(pred - target)
    loss = torch.where(diff < beta, 0.5 * diff * diff / beta,
                       diff - 0.5 * beta)
    return loss


Jiangmiao Pang's avatar
Jiangmiao Pang committed
18
19
20
@LOSSES.register_module
class SmoothL1Loss(nn.Module):

Kai Chen's avatar
Kai Chen committed
21
    def __init__(self, beta=1.0, reduction='mean', loss_weight=1.0):
Jiangmiao Pang's avatar
Jiangmiao Pang committed
22
23
        super(SmoothL1Loss, self).__init__()
        self.beta = beta
Kai Chen's avatar
Kai Chen committed
24
        self.reduction = reduction
Jiangmiao Pang's avatar
Jiangmiao Pang committed
25
26
        self.loss_weight = loss_weight

27
28
29
30
31
32
33
34
35
36
    def forward(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None,
                **kwargs):
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
Kai Chen's avatar
Kai Chen committed
37
38
39
40
41
        loss_bbox = self.loss_weight * smooth_l1_loss(
            pred,
            target,
            weight,
            beta=self.beta,
42
            reduction=reduction,
Kai Chen's avatar
Kai Chen committed
43
44
            avg_factor=avg_factor,
            **kwargs)
Jiangmiao Pang's avatar
Jiangmiao Pang committed
45
        return loss_bbox