chanelisa/objectdetectionhw
0
1# Loss functions2 3import torch4import torch.nn as nn5import torch.nn.functional as F6 7from utils.general import bbox_iou, bbox_alpha_iou, box_iou, box_giou, box_diou, box_ciou, xywh2xyxy8from utils.torch_utils import is_parallel9 10 11def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-59802844112 # return positive, negative label smoothing BCE targets13 return 1.0 - 0.5 * eps, 0.5 * eps14 15 16class BCEBlurWithLogitsLoss(nn.Module):17 # BCEwithLogitLoss() with reduced missing label effects.18 def __init__(self, alpha=0.05):19 super(BCEBlurWithLogitsLoss, self).__init__()20 self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss()21 self.alpha = alpha22 23 def forward(self, pred, true):24 loss = self.loss_fcn(pred, true)25 pred = torch.sigmoid(pred) # prob from logits26 dx = pred - true # reduce only missing label effects27 # dx = (pred - true).abs() # reduce missing label and false label effects28 alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4))29 loss *= alpha_factor30 return loss.mean()31 32 33class SigmoidBin(nn.Module):34 stride = None # strides computed during build35 export = False # onnx export36 37 def __init__(self, bin_count=10, min=0.0, max=1.0, reg_scale = 2.0, use_loss_regression=True, use_fw_regression=True, BCE_weight=1.0, smooth_eps=0.0):38 super(SigmoidBin, self).__init__()39 40 self.bin_count = bin_count41 self.length = bin_count + 142 self.min = min43 self.max = max44 self.scale = float(max - min)45 self.shift = self.scale / 2.046 47 self.use_loss_regression = use_loss_regression48 self.use_fw_regression = use_fw_regression49 self.reg_scale = reg_scale50 self.BCE_weight = BCE_weight51 52 start = min + (self.scale/2.0) / self.bin_count53 end = max - (self.scale/2.0) / self.bin_count54 step = self.scale / self.bin_count55 self.step = step56 #print(f" start = {start}, end = {end}, step = {step} ")57 58 bins = torch.range(start, end + 0.0001, step).float() 59 self.register_buffer('bins', bins) 60 61 62 self.cp = 1.0 - 0.5 * smooth_eps63 self.cn = 0.5 * smooth_eps64 65 self.BCEbins = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([BCE_weight]))66 self.MSELoss = nn.MSELoss()67 68 def get_length(self):69 return self.length70 71 def forward(self, pred):72 assert pred.shape[-1] == self.length, 'pred.shape[-1]=%d is not equal to self.length=%d' % (pred.shape[-1], self.length)73 74 pred_reg = (pred[..., 0] * self.reg_scale - self.reg_scale/2.0) * self.step75 pred_bin = pred[..., 1:(1+self.bin_count)]76 77 _, bin_idx = torch.max(pred_bin, dim=-1)78 bin_bias = self.bins[bin_idx]79 80 if self.use_fw_regression:81 result = pred_reg + bin_bias82 else:83 result = bin_bias84 result = result.clamp(min=self.min, max=self.max)85 86 return result87 88 89 def training_loss(self, pred, target):90 assert pred.shape[-1] == self.length, 'pred.shape[-1]=%d is not equal to self.length=%d' % (pred.shape[-1], self.length)91 assert pred.shape[0] == target.shape[0], 'pred.shape=%d is not equal to the target.shape=%d' % (pred.shape[0], target.shape[0])92 device = pred.device93 94 pred_reg = (pred[..., 0].sigmoid() * self.reg_scale - self.reg_scale/2.0) * self.step95 pred_bin = pred[..., 1:(1+self.bin_count)]96 97 diff_bin_target = torch.abs(target[..., None] - self.bins)98 _, bin_idx = torch.min(diff_bin_target, dim=-1)99 100 bin_bias = self.bins[bin_idx]101 bin_bias.requires_grad = False102 result = pred_reg + bin_bias103 104 target_bins = torch.full_like(pred_bin, self.cn, device=device) # targets105 n = pred.shape[0] 106 target_bins[range(n), bin_idx] = self.cp107 108 loss_bin = self.BCEbins(pred_bin, target_bins) # BCE109 110 if self.use_loss_regression:111 loss_regression = self.MSELoss(result, target) # MSE 112 loss = loss_bin + loss_regression113 else:114 loss = loss_bin115 116 out_result = result.clamp(min=self.min, max=self.max)117 118 return loss, out_result119 120 121class FocalLoss(nn.Module):122 # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)123 def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):124 super(FocalLoss, self).__init__()125 self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()126 self.gamma = gamma127 self.alpha = alpha128 self.reduction = loss_fcn.reduction129 self.loss_fcn.reduction = 'none' # required to apply FL to each element130 131 def forward(self, pred, true):132 loss = self.loss_fcn(pred, true)133 # p_t = torch.exp(-loss)134 # loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability135 136 # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py137 pred_prob = torch.sigmoid(pred) # prob from logits138 p_t = true * pred_prob + (1 - true) * (1 - pred_prob)139 alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)140 modulating_factor = (1.0 - p_t) ** self.gamma141 loss *= alpha_factor * modulating_factor142 143 if self.reduction == 'mean':144 return loss.mean()145 elif self.reduction == 'sum':146 return loss.sum()147 else: # 'none'148 return loss149 150 151class QFocalLoss(nn.Module):152 # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)153 def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):154 super(QFocalLoss, self).__init__()155 self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()156 self.gamma = gamma157 self.alpha = alpha158 self.reduction = loss_fcn.reduction159 self.loss_fcn.reduction = 'none' # required to apply FL to each element160 161 def forward(self, pred, true):162 loss = self.loss_fcn(pred, true)163 164 pred_prob = torch.sigmoid(pred) # prob from logits165 alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)166 modulating_factor = torch.abs(true - pred_prob) ** self.gamma167 loss *= alpha_factor * modulating_factor168 169 if self.reduction == 'mean':170 return loss.mean()171 elif self.reduction == 'sum':172 return loss.sum()173 else: # 'none'174 return loss175 176class RankSort(torch.autograd.Function):177 @staticmethod178 def forward(ctx, logits, targets, delta_RS=0.50, eps=1e-10): 179 180 classification_grads=torch.zeros(logits.shape).cuda()181 182 #Filter fg logits183 fg_labels = (targets > 0.)184 fg_logits = logits[fg_labels]185 fg_targets = targets[fg_labels]186 fg_num = len(fg_logits)187 188 #Do not use bg with scores less than minimum fg logit189 #since changing its score does not have an effect on precision190 threshold_logit = torch.min(fg_logits)-delta_RS191 relevant_bg_labels=((targets==0) & (logits>=threshold_logit))192 193 relevant_bg_logits = logits[relevant_bg_labels] 194 relevant_bg_grad=torch.zeros(len(relevant_bg_logits)).cuda()195 sorting_error=torch.zeros(fg_num).cuda()196 ranking_error=torch.zeros(fg_num).cuda()197 fg_grad=torch.zeros(fg_num).cuda()198 199 #sort the fg logits200 order=torch.argsort(fg_logits)201 #Loops over each positive following the order202 for ii in order:203 # Difference Transforms (x_ij)204 fg_relations=fg_logits-fg_logits[ii] 205 bg_relations=relevant_bg_logits-fg_logits[ii]206 207 if delta_RS > 0:208 fg_relations=torch.clamp(fg_relations/(2*delta_RS)+0.5,min=0,max=1)209 bg_relations=torch.clamp(bg_relations/(2*delta_RS)+0.5,min=0,max=1)210 else:211 fg_relations = (fg_relations >= 0).float()212 bg_relations = (bg_relations >= 0).float()213 214 # Rank of ii among pos and false positive number (bg with larger scores)215 rank_pos=torch.sum(fg_relations)216 FP_num=torch.sum(bg_relations)217 218 # Rank of ii among all examples219 rank=rank_pos+FP_num220 221 # Ranking error of example ii. target_ranking_error is always 0. (Eq. 7)222 ranking_error[ii]=FP_num/rank 223 224 # Current sorting error of example ii. (Eq. 7)225 current_sorting_error = torch.sum(fg_relations*(1-fg_targets))/rank_pos226 227 #Find examples in the target sorted order for example ii 228 iou_relations = (fg_targets >= fg_targets[ii])229 target_sorted_order = iou_relations * fg_relations230 231 #The rank of ii among positives in sorted order232 rank_pos_target = torch.sum(target_sorted_order)233 234 #Compute target sorting error. (Eq. 8)235 #Since target ranking error is 0, this is also total target error 236 target_sorting_error= torch.sum(target_sorted_order*(1-fg_targets))/rank_pos_target237 238 #Compute sorting error on example ii239 sorting_error[ii] = current_sorting_error - target_sorting_error240 241 #Identity Update for Ranking Error 242 if FP_num > eps:243 #For ii the update is the ranking error244 fg_grad[ii] -= ranking_error[ii]245 #For negatives, distribute error via ranking pmf (i.e. bg_relations/FP_num)246 relevant_bg_grad += (bg_relations*(ranking_error[ii]/FP_num))247 248 #Find the positives that are misranked (the cause of the error)249 #These are the ones with smaller IoU but larger logits250 missorted_examples = (~ iou_relations) * fg_relations251 252 #Denominotor of sorting pmf 253 sorting_pmf_denom = torch.sum(missorted_examples)254 255 #Identity Update for Sorting Error 256 if sorting_pmf_denom > eps:257 #For ii the update is the sorting error258 fg_grad[ii] -= sorting_error[ii]259 #For positives, distribute error via sorting pmf (i.e. missorted_examples/sorting_pmf_denom)260 fg_grad += (missorted_examples*(sorting_error[ii]/sorting_pmf_denom))261 262 #Normalize gradients by number of positives 263 classification_grads[fg_labels]= (fg_grad/fg_num)264 classification_grads[relevant_bg_labels]= (relevant_bg_grad/fg_num)265 266 ctx.save_for_backward(classification_grads)267 268 return ranking_error.mean(), sorting_error.mean()269 270 @staticmethod271 def backward(ctx, out_grad1, out_grad2):272 g1, =ctx.saved_tensors273 return g1*out_grad1, None, None, None274 275class aLRPLoss(torch.autograd.Function):276 @staticmethod277 def forward(ctx, logits, targets, regression_losses, delta=1., eps=1e-5): 278 classification_grads=torch.zeros(logits.shape).cuda()279 280 #Filter fg logits281 fg_labels = (targets == 1)282 fg_logits = logits[fg_labels]283 fg_num = len(fg_logits)284 285 #Do not use bg with scores less than minimum fg logit286 #since changing its score does not have an effect on precision287 threshold_logit = torch.min(fg_logits)-delta288 289 #Get valid bg logits290 relevant_bg_labels=((targets==0)&(logits>=threshold_logit))291 relevant_bg_logits=logits[relevant_bg_labels] 292 relevant_bg_grad=torch.zeros(len(relevant_bg_logits)).cuda()293 rank=torch.zeros(fg_num).cuda()294 prec=torch.zeros(fg_num).cuda()295 fg_grad=torch.zeros(fg_num).cuda()296 297 max_prec=0 298 #sort the fg logits299 order=torch.argsort(fg_logits)300 #Loops over each positive following the order301 for ii in order:302 #x_ij s as score differences with fgs303 fg_relations=fg_logits-fg_logits[ii] 304 #Apply piecewise linear function and determine relations with fgs305 fg_relations=torch.clamp(fg_relations/(2*delta)+0.5,min=0,max=1)306 #Discard i=j in the summation in rank_pos307 fg_relations[ii]=0308 309 #x_ij s as score differences with bgs310 bg_relations=relevant_bg_logits-fg_logits[ii]311 #Apply piecewise linear function and determine relations with bgs312 bg_relations=torch.clamp(bg_relations/(2*delta)+0.5,min=0,max=1)313 314 #Compute the rank of the example within fgs and number of bgs with larger scores315 rank_pos=1+torch.sum(fg_relations)316 FP_num=torch.sum(bg_relations)317 #Store the total since it is normalizer also for aLRP Regression error318 rank[ii]=rank_pos+FP_num319 320 #Compute precision for this example to compute classification loss 321 prec[ii]=rank_pos/rank[ii] 322 #For stability, set eps to a infinitesmall value (e.g. 1e-6), then compute grads323 if FP_num > eps: 324 fg_grad[ii] = -(torch.sum(fg_relations*regression_losses)+FP_num)/rank[ii]325 relevant_bg_grad += (bg_relations*(-fg_grad[ii]/FP_num)) 326 327 #aLRP with grad formulation fg gradient328 classification_grads[fg_labels]= fg_grad329 #aLRP with grad formulation bg gradient330 classification_grads[relevant_bg_labels]= relevant_bg_grad 331 332 classification_grads /= (fg_num)333 334 cls_loss=1-prec.mean()335 ctx.save_for_backward(classification_grads)336 337 return cls_loss, rank, order338 339 @staticmethod340 def backward(ctx, out_grad1, out_grad2, out_grad3):341 g1, =ctx.saved_tensors342 return g1*out_grad1, None, None, None, None343 344 345class APLoss(torch.autograd.Function):346 @staticmethod347 def forward(ctx, logits, targets, delta=1.): 348 classification_grads=torch.zeros(logits.shape).cuda()349 350 #Filter fg logits351 fg_labels = (targets == 1)352 fg_logits = logits[fg_labels]353 fg_num = len(fg_logits)354 355 #Do not use bg with scores less than minimum fg logit356 #since changing its score does not have an effect on precision357 threshold_logit = torch.min(fg_logits)-delta358 359 #Get valid bg logits360 relevant_bg_labels=((targets==0)&(logits>=threshold_logit))361 relevant_bg_logits=logits[relevant_bg_labels] 362 relevant_bg_grad=torch.zeros(len(relevant_bg_logits)).cuda()363 rank=torch.zeros(fg_num).cuda()364 prec=torch.zeros(fg_num).cuda()365 fg_grad=torch.zeros(fg_num).cuda()366 367 max_prec=0 368 #sort the fg logits369 order=torch.argsort(fg_logits)370 #Loops over each positive following the order371 for ii in order:372 #x_ij s as score differences with fgs373 fg_relations=fg_logits-fg_logits[ii] 374 #Apply piecewise linear function and determine relations with fgs375 fg_relations=torch.clamp(fg_relations/(2*delta)+0.5,min=0,max=1)376 #Discard i=j in the summation in rank_pos377 fg_relations[ii]=0378 379 #x_ij s as score differences with bgs380 bg_relations=relevant_bg_logits-fg_logits[ii]381 #Apply piecewise linear function and determine relations with bgs382 bg_relations=torch.clamp(bg_relations/(2*delta)+0.5,min=0,max=1)383 384 #Compute the rank of the example within fgs and number of bgs with larger scores385 rank_pos=1+torch.sum(fg_relations)386 FP_num=torch.sum(bg_relations)387 #Store the total since it is normalizer also for aLRP Regression error388 rank[ii]=rank_pos+FP_num389 390 #Compute precision for this example 391 current_prec=rank_pos/rank[ii]392 393 #Compute interpolated AP and store gradients for relevant bg examples394 if (max_prec<=current_prec):395 max_prec=current_prec396 relevant_bg_grad += (bg_relations/rank[ii])397 else:398 relevant_bg_grad += (bg_relations/rank[ii])*(((1-max_prec)/(1-current_prec)))399 400 #Store fg gradients401 fg_grad[ii]=-(1-max_prec)402 prec[ii]=max_prec 403 404 #aLRP with grad formulation fg gradient405 classification_grads[fg_labels]= fg_grad406 #aLRP with grad formulation bg gradient407 classification_grads[relevant_bg_labels]= relevant_bg_grad 408 409 classification_grads /= fg_num410 411 cls_loss=1-prec.mean()412 ctx.save_for_backward(classification_grads)413 414 return cls_loss415 416 @staticmethod417 def backward(ctx, out_grad1):418 g1, =ctx.saved_tensors419 return g1*out_grad1, None, None420 421 422class ComputeLoss:423 # Compute losses424 def __init__(self, model, autobalance=False):425 super(ComputeLoss, self).__init__()426 device = next(model.parameters()).device # get model device427 h = model.hyp # hyperparameters428 429 # Define criteria430 BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))431 BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))432 433 # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3434 self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets435 436 # Focal loss437 g = h['fl_gamma'] # focal loss gamma438 if g > 0:439 BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)440 441 det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module442 self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P7443 #self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.1, .05]) # P3-P7444 #self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.5, 0.4, .1]) # P3-P7445 self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index446 self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance447 for k in 'na', 'nc', 'nl', 'anchors':448 setattr(self, k, getattr(det, k))449 450 def __call__(self, p, targets): # predictions, targets, model451 device = targets.device452 lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)453 tcls, tbox, indices, anchors = self.build_targets(p, targets) # targets454 455 # Losses456 for i, pi in enumerate(p): # layer index, layer predictions457 b, a, gj, gi = indices[i] # image, anchor, gridy, gridx458 tobj = torch.zeros_like(pi[..., 0], device=device) # target obj459 460 n = b.shape[0] # number of targets461 if n:462 ps = pi[b, a, gj, gi] # prediction subset corresponding to targets463 464 # Regression465 pxy = ps[:, :2].sigmoid() * 2. - 0.5466 pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]467 pbox = torch.cat((pxy, pwh), 1) # predicted box468 iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target)469 lbox += (1.0 - iou).mean() # iou loss470 471 # Objectness472 tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio473 474 # Classification475 if self.nc > 1: # cls loss (only if multiple classes)476 t = torch.full_like(ps[:, 5:], self.cn, device=device) # targets477 t[range(n), tcls[i]] = self.cp478 #t[t==self.cp] = iou.detach().clamp(0).type(t.dtype)479 lcls += self.BCEcls(ps[:, 5:], t) # BCE480 481 # Append targets to text file482 # with open('targets.txt', 'a') as file:483 # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]484 485 obji = self.BCEobj(pi[..., 4], tobj)486 lobj += obji * self.balance[i] # obj loss487 if self.autobalance:488 self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()489 490 if self.autobalance:491 self.balance = [x / self.balance[self.ssi] for x in self.balance]492 lbox *= self.hyp['box']493 lobj *= self.hyp['obj']494 lcls *= self.hyp['cls']495 bs = tobj.shape[0] # batch size496 497 loss = lbox + lobj + lcls498 return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()499 500 def build_targets(self, p, targets):501 # Build targets for compute_loss(), input targets(image,class,x,y,w,h)502 na, nt = self.na, targets.shape[0] # number of anchors, targets503 tcls, tbox, indices, anch = [], [], [], []504 gain = torch.ones(7, device=targets.device).long() # normalized to gridspace gain505 ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)506 targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices507 508 g = 0.5 # bias509 off = torch.tensor([[0, 0],510 [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m511 # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm512 ], device=targets.device).float() * g # offsets513 514 for i in range(self.nl):515 anchors = self.anchors[i]516 gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain517 518 # Match targets to anchors519 t = targets * gain520 if nt:521 # Matches522 r = t[:, :, 4:6] / anchors[:, None] # wh ratio523 j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t'] # compare524 # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))525 t = t[j] # filter526 527 # Offsets528 gxy = t[:, 2:4] # grid xy529 gxi = gain[[2, 3]] - gxy # inverse530 j, k = ((gxy % 1. < g) & (gxy > 1.)).T531 l, m = ((gxi % 1. < g) & (gxi > 1.)).T532 j = torch.stack((torch.ones_like(j), j, k, l, m))533 t = t.repeat((5, 1, 1))[j]534 offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]535 else:536 t = targets[0]537 offsets = 0538 539 # Define540 b, c = t[:, :2].long().T # image, class541 gxy = t[:, 2:4] # grid xy542 gwh = t[:, 4:6] # grid wh543 gij = (gxy - offsets).long()544 gi, gj = gij.T # grid xy indices545 546 # Append547 a = t[:, 6].long() # anchor indices548 indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices549 tbox.append(torch.cat((gxy - gij, gwh), 1)) # box550 anch.append(anchors[a]) # anchors551 tcls.append(c) # class552 553 return tcls, tbox, indices, anch554 555 556class ComputeLossOTA:557 # Compute losses558 def __init__(self, model, autobalance=False):559 super(ComputeLossOTA, self).__init__()560 device = next(model.parameters()).device # get model device561 h = model.hyp # hyperparameters562 563 # Define criteria564 BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))565 BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))566 567 # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3568 self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets569 570 # Focal loss571 g = h['fl_gamma'] # focal loss gamma572 if g > 0:573 BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)574 575 det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module576 self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P7577 self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index578 self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance579 for k in 'na', 'nc', 'nl', 'anchors', 'stride':580 setattr(self, k, getattr(det, k))581 582 def __call__(self, p, targets, imgs): # predictions, targets, model 583 device = targets.device584 lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)585 bs, as_, gjs, gis, targets, anchors = self.build_targets(p, targets, imgs)586 pre_gen_gains = [torch.tensor(pp.shape, device=device)[[3, 2, 3, 2]] for pp in p] 587 588 589 # Losses590 for i, pi in enumerate(p): # layer index, layer predictions591 b, a, gj, gi = bs[i], as_[i], gjs[i], gis[i] # image, anchor, gridy, gridx592 tobj = torch.zeros_like(pi[..., 0], device=device) # target obj593 594 n = b.shape[0] # number of targets595 if n:596 ps = pi[b, a, gj, gi] # prediction subset corresponding to targets597 598 # Regression599 grid = torch.stack([gi, gj], dim=1)600 pxy = ps[:, :2].sigmoid() * 2. - 0.5601 #pxy = ps[:, :2].sigmoid() * 3. - 1.602 pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]603 pbox = torch.cat((pxy, pwh), 1) # predicted box604 selected_tbox = targets[i][:, 2:6] * pre_gen_gains[i]605 selected_tbox[:, :2] -= grid606 iou = bbox_iou(pbox.T, selected_tbox, x1y1x2y2=False, CIoU=True) # iou(prediction, target)607 lbox += (1.0 - iou).mean() # iou loss608 609 # Objectness610 tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio611 612 # Classification613 selected_tcls = targets[i][:, 1].long()614 if self.nc > 1: # cls loss (only if multiple classes)615 t = torch.full_like(ps[:, 5:], self.cn, device=device) # targets616 t[range(n), selected_tcls] = self.cp617 lcls += self.BCEcls(ps[:, 5:], t) # BCE618 619 # Append targets to text file620 # with open('targets.txt', 'a') as file:621 # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]622 623 obji = self.BCEobj(pi[..., 4], tobj)624 lobj += obji * self.balance[i] # obj loss625 if self.autobalance:626 self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()627 628 if self.autobalance:629 self.balance = [x / self.balance[self.ssi] for x in self.balance]630 lbox *= self.hyp['box']631 lobj *= self.hyp['obj']632 lcls *= self.hyp['cls']633 bs = tobj.shape[0] # batch size634 635 loss = lbox + lobj + lcls636 return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()637 638 def build_targets(self, p, targets, imgs):639 640 #indices, anch = self.find_positive(p, targets)641 indices, anch = self.find_3_positive(p, targets)642 #indices, anch = self.find_4_positive(p, targets)643 #indices, anch = self.find_5_positive(p, targets)644 #indices, anch = self.find_9_positive(p, targets)645 device = torch.device(targets.device)646 matching_bs = [[] for pp in p]647 matching_as = [[] for pp in p]648 matching_gjs = [[] for pp in p]649 matching_gis = [[] for pp in p]650 matching_targets = [[] for pp in p]651 matching_anchs = [[] for pp in p]652 653 nl = len(p) 654 655 for batch_idx in range(p[0].shape[0]):656 657 b_idx = targets[:, 0]==batch_idx658 this_target = targets[b_idx]659 if this_target.shape[0] == 0:660 continue661 662 txywh = this_target[:, 2:6] * imgs[batch_idx].shape[1]663 txyxy = xywh2xyxy(txywh)664 665 pxyxys = []666 p_cls = []667 p_obj = []668 from_which_layer = []669 all_b = []670 all_a = []671 all_gj = []672 all_gi = []673 all_anch = []674 675 for i, pi in enumerate(p):676 677 b, a, gj, gi = indices[i]678 idx = (b == batch_idx)679 b, a, gj, gi = b[idx], a[idx], gj[idx], gi[idx] 680 all_b.append(b)681 all_a.append(a)682 all_gj.append(gj)683 all_gi.append(gi)684 all_anch.append(anch[i][idx])685 from_which_layer.append((torch.ones(size=(len(b),)) * i).to(device))686 687 fg_pred = pi[b, a, gj, gi] 688 p_obj.append(fg_pred[:, 4:5])689 p_cls.append(fg_pred[:, 5:])690 691 grid = torch.stack([gi, gj], dim=1)692 pxy = (fg_pred[:, :2].sigmoid() * 2. - 0.5 + grid) * self.stride[i] #/ 8.693 #pxy = (fg_pred[:, :2].sigmoid() * 3. - 1. + grid) * self.stride[i]694 pwh = (fg_pred[:, 2:4].sigmoid() * 2) ** 2 * anch[i][idx] * self.stride[i] #/ 8.695 pxywh = torch.cat([pxy, pwh], dim=-1)696 pxyxy = xywh2xyxy(pxywh)697 pxyxys.append(pxyxy)698 699 pxyxys = torch.cat(pxyxys, dim=0)700 if pxyxys.shape[0] == 0:701 continue702 p_obj = torch.cat(p_obj, dim=0)703 p_cls = torch.cat(p_cls, dim=0)704 from_which_layer = torch.cat(from_which_layer, dim=0)705 all_b = torch.cat(all_b, dim=0)706 all_a = torch.cat(all_a, dim=0)707 all_gj = torch.cat(all_gj, dim=0)708 all_gi = torch.cat(all_gi, dim=0)709 all_anch = torch.cat(all_anch, dim=0)710 711 pair_wise_iou = box_iou(txyxy, pxyxys)712 713 pair_wise_iou_loss = -torch.log(pair_wise_iou + 1e-8)714 715 top_k, _ = torch.topk(pair_wise_iou, min(10, pair_wise_iou.shape[1]), dim=1)716 dynamic_ks = torch.clamp(top_k.sum(1).int(), min=1)717 718 gt_cls_per_image = (719 F.one_hot(this_target[:, 1].to(torch.int64), self.nc)720 .float()721 .unsqueeze(1)722 .repeat(1, pxyxys.shape[0], 1)723 )724 725 num_gt = this_target.shape[0]726 cls_preds_ = (727 p_cls.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()728 * p_obj.unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()729 )730 731 y = cls_preds_.sqrt_()732 pair_wise_cls_loss = F.binary_cross_entropy_with_logits(733 torch.log(y/(1-y)) , gt_cls_per_image, reduction="none"734 ).sum(-1)735 del cls_preds_736 737 cost = (738 pair_wise_cls_loss739 + 3.0 * pair_wise_iou_loss740 )741 742 matching_matrix = torch.zeros_like(cost, device=device)743 744 for gt_idx in range(num_gt):745 _, pos_idx = torch.topk(746 cost[gt_idx], k=dynamic_ks[gt_idx].item(), largest=False747 )748 matching_matrix[gt_idx][pos_idx] = 1.0749 750 del top_k, dynamic_ks751 anchor_matching_gt = matching_matrix.sum(0)752 if (anchor_matching_gt > 1).sum() > 0:753 _, cost_argmin = torch.min(cost[:, anchor_matching_gt > 1], dim=0)754 matching_matrix[:, anchor_matching_gt > 1] *= 0.0755 matching_matrix[cost_argmin, anchor_matching_gt > 1] = 1.0756 fg_mask_inboxes = (matching_matrix.sum(0) > 0.0).to(device)757 matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)758 759 from_which_layer = from_which_layer[fg_mask_inboxes]760 all_b = all_b[fg_mask_inboxes]761 all_a = all_a[fg_mask_inboxes]762 all_gj = all_gj[fg_mask_inboxes]763 all_gi = all_gi[fg_mask_inboxes]764 all_anch = all_anch[fg_mask_inboxes]765 766 this_target = this_target[matched_gt_inds]767 768 for i in range(nl):769 layer_idx = from_which_layer == i770 matching_bs[i].append(all_b[layer_idx])771 matching_as[i].append(all_a[layer_idx])772 matching_gjs[i].append(all_gj[layer_idx])773 matching_gis[i].append(all_gi[layer_idx])774 matching_targets[i].append(this_target[layer_idx])775 matching_anchs[i].append(all_anch[layer_idx])776 777 for i in range(nl):778 if matching_targets[i] != []:779 matching_bs[i] = torch.cat(matching_bs[i], dim=0)780 matching_as[i] = torch.cat(matching_as[i], dim=0)781 matching_gjs[i] = torch.cat(matching_gjs[i], dim=0)782 matching_gis[i] = torch.cat(matching_gis[i], dim=0)783 matching_targets[i] = torch.cat(matching_targets[i], dim=0)784 matching_anchs[i] = torch.cat(matching_anchs[i], dim=0)785 else:786 matching_bs[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)787 matching_as[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)788 matching_gjs[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)789 matching_gis[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)790 matching_targets[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)791 matching_anchs[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)792 793 return matching_bs, matching_as, matching_gjs, matching_gis, matching_targets, matching_anchs 794 795 def find_3_positive(self, p, targets):796 # Build targets for compute_loss(), input targets(image,class,x,y,w,h)797 na, nt = self.na, targets.shape[0] # number of anchors, targets798 indices, anch = [], []799 gain = torch.ones(7, device=targets.device).long() # normalized to gridspace gain800 ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)801 targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices802 803 g = 0.5 # bias804 off = torch.tensor([[0, 0],805 [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m806 # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm807 ], device=targets.device).float() * g # offsets808 809 for i in range(self.nl):810 anchors = self.anchors[i]811 gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain812 813 # Match targets to anchors814 t = targets * gain815 if nt:816 # Matches817 r = t[:, :, 4:6] / anchors[:, None] # wh ratio818 j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t'] # compare819 # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))820 t = t[j] # filter821 822 # Offsets823 gxy = t[:, 2:4] # grid xy824 gxi = gain[[2, 3]] - gxy # inverse825 j, k = ((gxy % 1. < g) & (gxy > 1.)).T826 l, m = ((gxi % 1. < g) & (gxi > 1.)).T827 j = torch.stack((torch.ones_like(j), j, k, l, m))828 t = t.repeat((5, 1, 1))[j]829 offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]830 else:831 t = targets[0]832 offsets = 0833 834 # Define835 b, c = t[:, :2].long().T # image, class836 gxy = t[:, 2:4] # grid xy837 gwh = t[:, 4:6] # grid wh838 gij = (gxy - offsets).long()839 gi, gj = gij.T # grid xy indices840 841 # Append842 a = t[:, 6].long() # anchor indices843 indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices844 anch.append(anchors[a]) # anchors845 846 return indices, anch847 848 849class ComputeLossBinOTA:850 # Compute losses851 def __init__(self, model, autobalance=False):852 super(ComputeLossBinOTA, self).__init__()853 device = next(model.parameters()).device # get model device854 h = model.hyp # hyperparameters855 856 # Define criteria857 BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))858 BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))859 #MSEangle = nn.MSELoss().to(device)860 861 # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3862 self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets863 864 # Focal loss865 g = h['fl_gamma'] # focal loss gamma866 if g > 0:867 BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)868 869 det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module870 self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P7871 self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index872 self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance873 for k in 'na', 'nc', 'nl', 'anchors', 'stride', 'bin_count':874 setattr(self, k, getattr(det, k))875 876 #xy_bin_sigmoid = SigmoidBin(bin_count=11, min=-0.5, max=1.5, use_loss_regression=False).to(device)877 wh_bin_sigmoid = SigmoidBin(bin_count=self.bin_count, min=0.0, max=4.0, use_loss_regression=False).to(device)878 #angle_bin_sigmoid = SigmoidBin(bin_count=31, min=-1.1, max=1.1, use_loss_regression=False).to(device)879 self.wh_bin_sigmoid = wh_bin_sigmoid880 881 def __call__(self, p, targets, imgs): # predictions, targets, model 882 device = targets.device883 lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)884 bs, as_, gjs, gis, targets, anchors = self.build_targets(p, targets, imgs)885 pre_gen_gains = [torch.tensor(pp.shape, device=device)[[3, 2, 3, 2]] for pp in p] 886 887 888 # Losses889 for i, pi in enumerate(p): # layer index, layer predictions890 b, a, gj, gi = bs[i], as_[i], gjs[i], gis[i] # image, anchor, gridy, gridx891 tobj = torch.zeros_like(pi[..., 0], device=device) # target obj892 893 obj_idx = self.wh_bin_sigmoid.get_length()*2 + 2 # x,y, w-bce, h-bce # xy_bin_sigmoid.get_length()*2894 895 n = b.shape[0] # number of targets896 if n:897 ps = pi[b, a, gj, gi] # prediction subset corresponding to targets898 899 # Regression900 grid = torch.stack([gi, gj], dim=1)901 selected_tbox = targets[i][:, 2:6] * pre_gen_gains[i]902 selected_tbox[:, :2] -= grid903 904 #pxy = ps[:, :2].sigmoid() * 2. - 0.5905 ##pxy = ps[:, :2].sigmoid() * 3. - 1.906 #pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]907 #pbox = torch.cat((pxy, pwh), 1) # predicted box908 909 #x_loss, px = xy_bin_sigmoid.training_loss(ps[..., 0:12], tbox[i][..., 0])910 #y_loss, py = xy_bin_sigmoid.training_loss(ps[..., 12:24], tbox[i][..., 1])911 w_loss, pw = self.wh_bin_sigmoid.training_loss(ps[..., 2:(3+self.bin_count)], selected_tbox[..., 2] / anchors[i][..., 0])912 h_loss, ph = self.wh_bin_sigmoid.training_loss(ps[..., (3+self.bin_count):obj_idx], selected_tbox[..., 3] / anchors[i][..., 1])913 914 pw *= anchors[i][..., 0]915 ph *= anchors[i][..., 1]916 917 px = ps[:, 0].sigmoid() * 2. - 0.5918 py = ps[:, 1].sigmoid() * 2. - 0.5919 920 lbox += w_loss + h_loss # + x_loss + y_loss921 922 #print(f"\n px = {px.shape}, py = {py.shape}, pw = {pw.shape}, ph = {ph.shape} \n")923 924 pbox = torch.cat((px.unsqueeze(1), py.unsqueeze(1), pw.unsqueeze(1), ph.unsqueeze(1)), 1).to(device) # predicted box925 926 927 928 929 iou = bbox_iou(pbox.T, selected_tbox, x1y1x2y2=False, CIoU=True) # iou(prediction, target)930 lbox += (1.0 - iou).mean() # iou loss931 932 # Objectness933 tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio934 935 # Classification936 selected_tcls = targets[i][:, 1].long()937 if self.nc > 1: # cls loss (only if multiple classes)938 t = torch.full_like(ps[:, (1+obj_idx):], self.cn, device=device) # targets939 t[range(n), selected_tcls] = self.cp940 lcls += self.BCEcls(ps[:, (1+obj_idx):], t) # BCE941 942 # Append targets to text file943 # with open('targets.txt', 'a') as file:944 # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]945 946 obji = self.BCEobj(pi[..., obj_idx], tobj)947 lobj += obji * self.balance[i] # obj loss948 if self.autobalance:949 self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()950 951 if self.autobalance:952 self.balance = [x / self.balance[self.ssi] for x in self.balance]953 lbox *= self.hyp['box']954 lobj *= self.hyp['obj']955 lcls *= self.hyp['cls']956 bs = tobj.shape[0] # batch size957 958 loss = lbox + lobj + lcls959 return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()960 961 def build_targets(self, p, targets, imgs):962 963 #indices, anch = self.find_positive(p, targets)964 indices, anch = self.find_3_positive(p, targets)965 #indices, anch = self.find_4_positive(p, targets)966 #indices, anch = self.find_5_positive(p, targets)967 #indices, anch = self.find_9_positive(p, targets)968 969 matching_bs = [[] for pp in p]970 matching_as = [[] for pp in p]971 matching_gjs = [[] for pp in p]972 matching_gis = [[] for pp in p]973 matching_targets = [[] for pp in p]974 matching_anchs = [[] for pp in p]975 976 nl = len(p) 977 978 for batch_idx in range(p[0].shape[0]):979 980 b_idx = targets[:, 0]==batch_idx981 this_target = targets[b_idx]982 if this_target.shape[0] == 0:983 continue984 985 txywh = this_target[:, 2:6] * imgs[batch_idx].shape[1]986 txyxy = xywh2xyxy(txywh)987 988 pxyxys = []989 p_cls = []990 p_obj = []991 from_which_layer = []992 all_b = []993 all_a = []994 all_gj = []995 all_gi = []996 all_anch = []997 998 for i, pi in enumerate(p):999 1000 obj_idx = self.wh_bin_sigmoid.get_length()*2 + 21001 1002 b, a, gj, gi = indices[i]1003 idx = (b == batch_idx)1004 b, a, gj, gi = b[idx], a[idx], gj[idx], gi[idx] 1005 all_b.append(b)1006 all_a.append(a)1007 all_gj.append(gj)1008 all_gi.append(gi)1009 all_anch.append(anch[i][idx])1010 from_which_layer.append(torch.ones(size=(len(b),)) * i)1011 1012 fg_pred = pi[b, a, gj, gi] 1013 p_obj.append(fg_pred[:, obj_idx:(obj_idx+1)])1014 p_cls.append(fg_pred[:, (obj_idx+1):])1015 1016 grid = torch.stack([gi, gj], dim=1)1017 pxy = (fg_pred[:, :2].sigmoid() * 2. - 0.5 + grid) * self.stride[i] #/ 8.1018 #pwh = (fg_pred[:, 2:4].sigmoid() * 2) ** 2 * anch[i][idx] * self.stride[i] #/ 8.1019 pw = self.wh_bin_sigmoid.forward(fg_pred[..., 2:(3+self.bin_count)].sigmoid()) * anch[i][idx][:, 0] * self.stride[i]1020 ph = self.wh_bin_sigmoid.forward(fg_pred[..., (3+self.bin_count):obj_idx].sigmoid()) * anch[i][idx][:, 1] * self.stride[i]1021 1022 pxywh = torch.cat([pxy, pw.unsqueeze(1), ph.unsqueeze(1)], dim=-1)1023 pxyxy = xywh2xyxy(pxywh)1024 pxyxys.append(pxyxy)1025 1026 pxyxys = torch.cat(pxyxys, dim=0)1027 if pxyxys.shape[0] == 0:1028 continue1029 p_obj = torch.cat(p_obj, dim=0)1030 p_cls = torch.cat(p_cls, dim=0)1031 from_which_layer = torch.cat(from_which_layer, dim=0)1032 all_b = torch.cat(all_b, dim=0)1033 all_a = torch.cat(all_a, dim=0)1034 all_gj = torch.cat(all_gj, dim=0)1035 all_gi = torch.cat(all_gi, dim=0)1036 all_anch = torch.cat(all_anch, dim=0)1037 1038 pair_wise_iou = box_iou(txyxy, pxyxys)1039 1040 pair_wise_iou_loss = -torch.log(pair_wise_iou + 1e-8)1041 1042 top_k, _ = torch.topk(pair_wise_iou, min(10, pair_wise_iou.shape[1]), dim=1)1043 dynamic_ks = torch.clamp(top_k.sum(1).int(), min=1)1044 1045 gt_cls_per_image = (1046 F.one_hot(this_target[:, 1].to(torch.int64), self.nc)1047 .float()1048 .unsqueeze(1)1049 .repeat(1, pxyxys.shape[0], 1)1050 )1051 1052 num_gt = this_target.shape[0] 1053 cls_preds_ = (1054 p_cls.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()1055 * p_obj.unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()1056 )1057 1058 y = cls_preds_.sqrt_()1059 pair_wise_cls_loss = F.binary_cross_entropy_with_logits(1060 torch.log(y/(1-y)) , gt_cls_per_image, reduction="none"1061 ).sum(-1)1062 del cls_preds_1063 1064 cost = (1065 pair_wise_cls_loss1066 + 3.0 * pair_wise_iou_loss1067 )1068 1069 matching_matrix = torch.zeros_like(cost)1070 1071 for gt_idx in range(num_gt):1072 _, pos_idx = torch.topk(1073 cost[gt_idx], k=dynamic_ks[gt_idx].item(), largest=False1074 )1075 matching_matrix[gt_idx][pos_idx] = 1.01076 1077 del top_k, dynamic_ks1078 anchor_matching_gt = matching_matrix.sum(0)1079 if (anchor_matching_gt > 1).sum() > 0:1080 _, cost_argmin = torch.min(cost[:, anchor_matching_gt > 1], dim=0)1081 matching_matrix[:, anchor_matching_gt > 1] *= 0.01082 matching_matrix[cost_argmin, anchor_matching_gt > 1] = 1.01083 fg_mask_inboxes = matching_matrix.sum(0) > 0.01084 matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)1085 1086 from_which_layer = from_which_layer[fg_mask_inboxes]1087 all_b = all_b[fg_mask_inboxes]1088 all_a = all_a[fg_mask_inboxes]1089 all_gj = all_gj[fg_mask_inboxes]1090 all_gi = all_gi[fg_mask_inboxes]1091 all_anch = all_anch[fg_mask_inboxes]1092 1093 this_target = this_target[matched_gt_inds]1094 1095 for i in range(nl):1096 layer_idx = from_which_layer == i1097 matching_bs[i].append(all_b[layer_idx])1098 matching_as[i].append(all_a[layer_idx])1099 matching_gjs[i].append(all_gj[layer_idx])1100 matching_gis[i].append(all_gi[layer_idx])1101 matching_targets[i].append(this_target[layer_idx])1102 matching_anchs[i].append(all_anch[layer_idx])1103 1104 for i in range(nl):1105 if matching_targets[i] != []:1106 matching_bs[i] = torch.cat(matching_bs[i], dim=0)1107 matching_as[i] = torch.cat(matching_as[i], dim=0)1108 matching_gjs[i] = torch.cat(matching_gjs[i], dim=0)1109 matching_gis[i] = torch.cat(matching_gis[i], dim=0)1110 matching_targets[i] = torch.cat(matching_targets[i], dim=0)1111 matching_anchs[i] = torch.cat(matching_anchs[i], dim=0)1112 else:1113 matching_bs[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)1114 matching_as[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)1115 matching_gjs[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)1116 matching_gis[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)1117 matching_targets[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)1118 matching_anchs[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)1119 1120 return matching_bs, matching_as, matching_gjs, matching_gis, matching_targets, matching_anchs 1121 1122 def find_3_positive(self, p, targets):1123 # Build targets for compute_loss(), input targets(image,class,x,y,w,h)1124 na, nt = self.na, targets.shape[0] # number of anchors, targets1125 indices, anch = [], []1126 gain = torch.ones(7, device=targets.device).long() # normalized to gridspace gain1127 ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)1128 targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices1129 1130 g = 0.5 # bias1131 off = torch.tensor([[0, 0],1132 [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m1133 # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm1134 ], device=targets.device).float() * g # offsets1135 1136 for i in range(self.nl):1137 anchors = self.anchors[i]1138 gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain1139 1140 # Match targets to anchors1141 t = targets * gain1142 if nt:1143 # Matches1144 r = t[:, :, 4:6] / anchors[:, None] # wh ratio1145 j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t'] # compare1146 # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))1147 t = t[j] # filter1148 1149 # Offsets1150 gxy = t[:, 2:4] # grid xy1151 gxi = gain[[2, 3]] - gxy # inverse1152 j, k = ((gxy % 1. < g) & (gxy > 1.)).T1153 l, m = ((gxi % 1. < g) & (gxi > 1.)).T1154 j = torch.stack((torch.ones_like(j), j, k, l, m))1155 t = t.repeat((5, 1, 1))[j]1156 offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]1157 else:1158 t = targets[0]1159 offsets = 01160 1161 # Define1162 b, c = t[:, :2].long().T # image, class1163 gxy = t[:, 2:4] # grid xy1164 gwh = t[:, 4:6] # grid wh1165 gij = (gxy - offsets).long()1166 gi, gj = gij.T # grid xy indices1167 1168 # Append1169 a = t[:, 6].long() # anchor indices1170 indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices1171 anch.append(anchors[a]) # anchors1172 1173 return indices, anch1174 1175 1176class ComputeLossAuxOTA:1177 # Compute losses1178 def __init__(self, model, autobalance=False):1179 super(ComputeLossAuxOTA, self).__init__()1180 device = next(model.parameters()).device # get model device1181 h = model.hyp # hyperparameters1182 1183 # Define criteria1184 BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))1185 BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))1186 1187 # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 31188 self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets1189 1190 # Focal loss1191 g = h['fl_gamma'] # focal loss gamma1192 if g > 0:1193 BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)1194 1195 det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module1196 self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P71197 self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index1198 self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance1199 for k in 'na', 'nc', 'nl', 'anchors', 'stride':1200 setattr(self, k, getattr(det, k))