mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

Update objectness IoU sort (#3786)

Browse Source
This commit is contained in:
Glenn Jocher 2021-06-26 14:52:18 +02:00 committed by GitHub
parent 157aa2f886
commit 8035b61682
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
utils/loss.py
View File

@ -89,6 +89,7 @@ class ComputeLoss:
# Compute losses # Compute losses
def __init__(self, model, autobalance=False): def __init__(self, model, autobalance=False):
super(ComputeLoss, self).__init__() super(ComputeLoss, self).__init__()
self.sort_obj_iou = False
device = next(model.parameters()).device # get model device device = next(model.parameters()).device # get model device
h = model.hyp # hyperparameters h = model.hyp # hyperparameters
@ -134,8 +135,9 @@ class ComputeLoss:
# Objectness # Objectness
score_iou = iou.detach().clamp(0).type(tobj.dtype) score_iou = iou.detach().clamp(0).type(tobj.dtype)
sort_id = torch.argsort(score_iou) if self.sort_obj_iou:
b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[sort_id], gi[sort_id], score_iou[sort_id] 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 tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # iou ratio
# Classification # Classification