Update objectness IoU sort (#3786)

ultralytics · Jun 26, 2021 · 8035b61 · 8035b61
1 parent 157aa2f
commit 8035b61
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/utils/loss.py b/utils/loss.py
@@ -89,6 +89,7 @@ class ComputeLoss:
     # Compute losses
     def __init__(self, model, autobalance=False):
         super(ComputeLoss, self).__init__()
+        self.sort_obj_iou = False
         device = next(model.parameters()).device  # get model device
         h = model.hyp  # hyperparameters
 
@@ -134,8 +135,9 @@ def __call__(self, p, targets):  # predictions, targets, model
 
                 # Objectness
                 score_iou = iou.detach().clamp(0).type(tobj.dtype)
-                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]
+                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
 
                 # Classification