Update loss.py with if self.gr < 1: (#7087)

* Update loss.py with `if self.gr < 1:`

* Update loss.py
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-03-21 18:35:36 +01:00
1 parent 9b771a3e71
commit 6f128031d0
1 file changed
+6 -4
+6 -4
View File
@@ -139,11 +139,13 @@ class ComputeLoss:
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)