glenn-jocher banHuang commited on
Commit
157aa2f
·
unverified ·
1 Parent(s): 9dc5d35

Objectness IoU Sort (#3610)

Browse files

Co-authored-by: U-LAPTOP-5N89P8V7\banhu <[email protected]>

Files changed (1) hide show
  1. utils/loss.py +4 -1
utils/loss.py CHANGED
@@ -133,7 +133,10 @@ class ComputeLoss:
133
  lbox += (1.0 - iou).mean() # iou loss
134
 
135
  # Objectness
136
- tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio
 
 
 
137
 
138
  # Classification
139
  if self.nc > 1: # cls loss (only if multiple classes)
 
133
  lbox += (1.0 - iou).mean() # iou loss
134
 
135
  # Objectness
136
+ score_iou = iou.detach().clamp(0).type(tobj.dtype)
137
+ sort_id = torch.argsort(score_iou)
138
+ b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[sort_id], gi[sort_id], score_iou[sort_id]
139
+ tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # iou ratio
140
 
141
  # Classification
142
  if self.nc > 1: # cls loss (only if multiple classes)