Skip to content

Commit

Permalink
Update loss.py with if self.gr < 1: (#7087)
Browse files Browse the repository at this point in the history
* Update loss.py with `if self.gr < 1:`

* Update loss.py
  • Loading branch information
glenn-jocher committed Mar 21, 2022
1 parent 9b771a3 commit 6f12803
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions utils/loss.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,11 +139,13 @@ def __call__(self, p, targets): # predictions, targets
lbox += (1.0 - iou).mean() # iou loss

# Objectness
score_iou = iou.detach().clamp(0).type(tobj.dtype)
iou = iou.detach().clamp(0).type(tobj.dtype)
if self.sort_obj_iou:
sort_id = torch.argsort(score_iou)
b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[sort_id], gi[sort_id], score_iou[sort_id]
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # iou ratio
j = iou.argsort()
b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
if self.gr < 1:
iou = (1.0 - self.gr) + self.gr * iou
tobj[b, a, gj, gi] = iou # iou ratio

# Classification
if self.nc > 1: # cls loss (only if multiple classes)
Expand Down

0 comments on commit 6f12803

Please sign in to comment.