From 6b7681d6e9838f06cc43b5bab6d02d3986ccbc0c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 26 Jun 2021 14:48:18 +0200 Subject: [PATCH] Update objectness IoU sort --- utils/loss.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/utils/loss.py b/utils/loss.py index 576c1c79e6f8..d4c261a5cc97 100644 --- 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