Objectness IoU Sort (#3610)

Co-authored-by: U-LAPTOP-5N89P8V7\banhu <ban.huang@foxmail.com>
This commit is contained in:
Glenn Jocher 2021-06-26 14:45:53 +02:00 committed by GitHub
parent 9dc5d35fce
commit 157aa2f886
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions

View File

@ -133,7 +133,10 @@ class ComputeLoss:
lbox += (1.0 - iou).mean() # iou loss lbox += (1.0 - iou).mean() # iou loss
# Objectness # Objectness
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio 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]
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # iou ratio
# Classification # Classification
if self.nc > 1: # cls loss (only if multiple classes) if self.nc > 1: # cls loss (only if multiple classes)