31def cross_entropy_loss(pred_class_logits, gt_classes, eps, alpha=0.2):
32 num_classes = pred_class_logits.size(1)
33
34 if eps >= 0:
35 smooth_param = eps
36 else:
37
38 soft_label = F.softmax(pred_class_logits, dim=1)
39 smooth_param = alpha * soft_label[torch.arange(soft_label.size(0)), gt_classes].unsqueeze(1)
40
41 log_probs = F.log_softmax(pred_class_logits, dim=1)
42 with torch.no_grad():
43 targets = torch.ones_like(log_probs)
44 targets *= smooth_param / (num_classes - 1)
45 targets.scatter_(1, gt_classes.data.unsqueeze(1), (1 - smooth_param))
46
47 loss = (-targets * log_probs).sum(dim=1)
48
49 """
50 # confidence penalty
51 conf_penalty = 0.3
52 probs = F.softmax(pred_class_logits, dim=1)
53 entropy = torch.sum(-probs * log_probs, dim=1)
54 loss = torch.clamp_min(loss - conf_penalty * entropy, min=0.)
55 """
56
57 with torch.no_grad():
58 non_zero_cnt = max(loss.nonzero(as_tuple=False).size(0), 1)
59
60 loss = loss.sum() / non_zero_cnt
61
62 return loss