File size: 4,026 Bytes
93e84b8 |
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 |
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from functools import partial
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class MultiCrossEntropyLoss(nn.Module):
def __init__(self, focal=False, weight=None, reduce=True):
super(MultiCrossEntropyLoss, self).__init__()
self.num_classes = 23
self.focal = focal
self.weight= weight
self.reduce = reduce
self.gamma_ = torch.zeros(self.num_classes).to(device) + 0.025
self.gamma_f = 0.05
self.register_buffer('pos_grad', torch.zeros(self.num_classes-1).to(device))
self.register_buffer('neg_grad', torch.zeros(self.num_classes-1).to(device))
self.register_buffer('pos_neg', torch.ones(self.num_classes-1).to(device))
def forward(self, input, target):
target_sum = torch.sum(target, dim=1)
target_div = torch.where(target_sum != 0, target_sum, torch.ones_like(target_sum)).unsqueeze(1)
target = target / target_div
logsoftmax = nn.LogSoftmax(dim=1).to(input.device)
gamma = self.gamma_.clone()
gamma[:-1] = gamma[:-1] + self.gamma_f * (1 - self.pos_neg)
if not self.focal:
if self.weight is None:
output = torch.sum(-target * logsoftmax(input), 1)
else:
output = torch.sum(-target * logsoftmax(input) / self.weight, 1)
else:
softmax = nn.Softmax(dim=1).to(input.device)
p = softmax(input)
output = torch.sum(-target * (1 - p)**gamma * logsoftmax(input), 1)
if self.reduce:
return torch.mean(output)
else:
return output
def map_func(self, x, s):
min_val = torch.min(x)
max_val = torch.max(x)
mu = torch.mean(x)
x = (x - min_val) / (max_val - min_val)
return 1 / (1 + torch.exp(-s * (x - mu)))
def collect_grad(self, target, grad):
grad = torch.abs(grad.reshape(-1, grad.shape[-1])).to(device)
target = target.reshape(-1, target.shape[-1]).to(device)
pos_grad = torch.sum(grad * target, dim=0)[:-1]
neg_grad = torch.sum(grad * (1 - target), dim=0)[:-1]
self.pos_grad += pos_grad
self.neg_grad += neg_grad
self.pos_neg = torch.clamp(self.pos_grad / (self.neg_grad + 1e-10), min=0, max=1)
self.pos_neg = self.map_func(self.pos_neg, 1)
def cls_loss_func(y, output, use_focal=False, weight=None, reduce=True):
input_size = y.size()
y = y.float().to(device)
if weight is not None:
weight = weight.to(device)
loss_func = MultiCrossEntropyLoss(focal=True, weight=weight, reduce=reduce)
y = y.reshape(-1, y.size(-1))
output = output.reshape(-1, output.size(-1))
loss = loss_func(output, y)
if not reduce:
loss = loss.reshape(input_size[:-1])
return loss
def cls_loss_func_(loss_func, y, output, use_focal=False, weight=None, reduce=True):
input_size = y.size()
y = y.float().to(device)
if weight is not None:
weight = weight.to(device)
y = y.reshape(-1, y.size(-1))
output = output.reshape(-1, output.size(-1))
loss = loss_func(output, y)
if not reduce:
loss = loss.reshape(input_size[:-1])
return loss
def regress_loss_func(y, output):
y = y.float().to(device)
y = y.reshape(-1, y.size(-1))
output = output.reshape(-1, output.size(-1))
bgmask = y[:, 1] < -1e2
fg_logits = output[~bgmask]
bg_logits = output[bgmask]
fg_target = y[~bgmask]
bg_target = y[bgmask]
loss = nn.functional.l1_loss(fg_logits, fg_target)
if loss.isnan():
return torch.tensor([0.0], requires_grad=True).to(device)
return loss
def suppress_loss_func(y, output):
y = y.float().to(device)
y = y.reshape(-1, y.size(-1))
output = output.reshape(-1, output.size(-1))
loss = nn.functional.binary_cross_entropy(output, y)
return loss
|