Update objectness IoU sort (ultralytics#3786)

fhkiel-mlaip · Jun 26, 2021 · 7bfb728 · 7bfb728
1 parent 77cfd67
commit 7bfb728
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