|
|
@@ -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) |