losses.py 3.91 KB
Newer Older
Kai Chen's avatar
Kai Chen committed
1
# TODO merge naive and weighted loss.
pangjm's avatar
pangjm committed
2
3
4
5
import torch
import torch.nn.functional as F


Kai Chen's avatar
Kai Chen committed
6
7
8
9
10
def weighted_nll_loss(pred, label, weight, avg_factor=None):
    if avg_factor is None:
        avg_factor = max(torch.sum(weight > 0).float().item(), 1.)
    raw = F.nll_loss(pred, label, reduction='none')
    return torch.sum(raw * weight)[None] / avg_factor
pangjm's avatar
pangjm committed
11
12


13
def weighted_cross_entropy(pred, label, weight, avg_factor=None, reduce=True):
Kai Chen's avatar
Kai Chen committed
14
15
16
    if avg_factor is None:
        avg_factor = max(torch.sum(weight > 0).float().item(), 1.)
    raw = F.cross_entropy(pred, label, reduction='none')
yhcao6's avatar
rename  
yhcao6 committed
17
    if reduce:
yhcao6's avatar
yhcao6 committed
18
        return torch.sum(raw * weight)[None] / avg_factor
yhcao6's avatar
yhcao6 committed
19
    else:
yhcao6's avatar
yhcao6 committed
20
        return raw * weight / avg_factor
pangjm's avatar
pangjm committed
21
22


Kai Chen's avatar
Kai Chen committed
23
24
25
def weighted_binary_cross_entropy(pred, label, weight, avg_factor=None):
    if avg_factor is None:
        avg_factor = max(torch.sum(weight > 0).float().item(), 1.)
pangjm's avatar
pangjm committed
26
27
    return F.binary_cross_entropy_with_logits(
        pred, label.float(), weight.float(),
Kai Chen's avatar
Kai Chen committed
28
        reduction='sum')[None] / avg_factor
pangjm's avatar
pangjm committed
29
30
31
32
33
34
35


def sigmoid_focal_loss(pred,
                       target,
                       weight,
                       gamma=2.0,
                       alpha=0.25,
Kai Chen's avatar
Kai Chen committed
36
                       reduction='mean'):
pangjm's avatar
pangjm committed
37
    pred_sigmoid = pred.sigmoid()
38
    target = target.type_as(pred)
pangjm's avatar
pangjm committed
39
40
41
    pt = (1 - pred_sigmoid) * target + pred_sigmoid * (1 - target)
    weight = (alpha * target + (1 - alpha) * (1 - target)) * weight
    weight = weight * pt.pow(gamma)
Kai Chen's avatar
Kai Chen committed
42
43
    loss = F.binary_cross_entropy_with_logits(
        pred, target, reduction='none') * weight
Kai Chen's avatar
Kai Chen committed
44
45
46
47
48
49
50
51
    reduction_enum = F._Reduction.get_enum(reduction)
    # none: 0, mean:1, sum: 2
    if reduction_enum == 0:
        return loss
    elif reduction_enum == 1:
        return loss.mean()
    elif reduction_enum == 2:
        return loss.sum()
pangjm's avatar
pangjm committed
52
53
54
55
56
57
58


def weighted_sigmoid_focal_loss(pred,
                                target,
                                weight,
                                gamma=2.0,
                                alpha=0.25,
Kai Chen's avatar
Kai Chen committed
59
                                avg_factor=None,
pangjm's avatar
pangjm committed
60
                                num_classes=80):
Kai Chen's avatar
Kai Chen committed
61
62
    if avg_factor is None:
        avg_factor = torch.sum(weight > 0).float().item() / num_classes + 1e-6
pangjm's avatar
pangjm committed
63
64
    return sigmoid_focal_loss(
        pred, target, weight, gamma=gamma, alpha=alpha,
Kai Chen's avatar
Kai Chen committed
65
        reduction='sum')[None] / avg_factor
pangjm's avatar
pangjm committed
66
67
68
69
70
71
72


def mask_cross_entropy(pred, target, label):
    num_rois = pred.size()[0]
    inds = torch.arange(0, num_rois, dtype=torch.long, device=pred.device)
    pred_slice = pred[inds, label].squeeze(1)
    return F.binary_cross_entropy_with_logits(
Kai Chen's avatar
Kai Chen committed
73
        pred_slice, target, reduction='mean')[None]
pangjm's avatar
pangjm committed
74
75


Kai Chen's avatar
Kai Chen committed
76
def smooth_l1_loss(pred, target, beta=1.0, reduction='mean'):
pangjm's avatar
pangjm committed
77
78
79
80
81
    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)
Kai Chen's avatar
Kai Chen committed
82
83
84
    reduction_enum = F._Reduction.get_enum(reduction)
    # none: 0, mean:1, sum: 2
    if reduction_enum == 0:
Kai Chen's avatar
Kai Chen committed
85
        return loss
Kai Chen's avatar
Kai Chen committed
86
    elif reduction_enum == 1:
Kai Chen's avatar
Kai Chen committed
87
        return loss.sum() / pred.numel()
Kai Chen's avatar
Kai Chen committed
88
    elif reduction_enum == 2:
Kai Chen's avatar
Kai Chen committed
89
90
91
92
93
94
95
96
        return loss.sum()


def weighted_smoothl1(pred, target, weight, beta=1.0, avg_factor=None):
    if avg_factor is None:
        avg_factor = torch.sum(weight > 0).float().item() / 4 + 1e-6
    loss = smooth_l1_loss(pred, target, beta, reduction='none')
    return torch.sum(loss * weight)[None] / avg_factor
pangjm's avatar
pangjm committed
97
98
99
100
101
102


def accuracy(pred, target, topk=1):
    if isinstance(topk, int):
        topk = (topk, )
        return_single = True
Kai Chen's avatar
Kai Chen committed
103
104
    else:
        return_single = False
pangjm's avatar
pangjm committed
105
106
107
108
109
110
111
112
113
114
115

    maxk = max(topk)
    _, pred_label = pred.topk(maxk, 1, True, True)
    pred_label = pred_label.t()
    correct = pred_label.eq(target.view(1, -1).expand_as(pred_label))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / pred.size(0)))
    return res[0] if return_single else res