From 6f128031d073754ee8ed6b6a85ecb6c0619cd0a7 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 21 Mar 2022 18:35:36 +0100 Subject: [PATCH] Update loss.py with `if self.gr < 1:` (#7087) * Update loss.py with `if self.gr < 1:` * Update loss.py --- utils/loss.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/utils/loss.py b/utils/loss.py index 0f01378..b49cc7f 100644 --- a/utils/loss.py +++ b/utils/loss.py @@ -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)