From 8035b61682cba7b10be24b0ab35cc0295f14d6cd Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 26 Jun 2021 14:52:18 +0200 Subject: [PATCH] Update objectness IoU sort (#3786) --- 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