Update objectness IoU sort (#3786)
Browse files- utils/loss.py +4 -2
utils/loss.py
CHANGED
@@ -89,6 +89,7 @@ class ComputeLoss:
|
|
89 |
# Compute losses
|
90 |
def __init__(self, model, autobalance=False):
|
91 |
super(ComputeLoss, self).__init__()
|
|
|
92 |
device = next(model.parameters()).device # get model device
|
93 |
h = model.hyp # hyperparameters
|
94 |
|
@@ -134,8 +135,9 @@ class ComputeLoss:
|
|
134 |
|
135 |
# Objectness
|
136 |
score_iou = iou.detach().clamp(0).type(tobj.dtype)
|
137 |
-
|
138 |
-
|
|
|
139 |
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # iou ratio
|
140 |
|
141 |
# Classification
|
|
|
89 |
# Compute losses
|
90 |
def __init__(self, model, autobalance=False):
|
91 |
super(ComputeLoss, self).__init__()
|
92 |
+
self.sort_obj_iou = False
|
93 |
device = next(model.parameters()).device # get model device
|
94 |
h = model.hyp # hyperparameters
|
95 |
|
|
|
135 |
|
136 |
# Objectness
|
137 |
score_iou = iou.detach().clamp(0).type(tobj.dtype)
|
138 |
+
if self.sort_obj_iou:
|
139 |
+
sort_id = torch.argsort(score_iou)
|
140 |
+
b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[sort_id], gi[sort_id], score_iou[sort_id]
|
141 |
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # iou ratio
|
142 |
|
143 |
# Classification
|