glenn-jocher commited on
Commit
8035b61
1 Parent(s): 157aa2f

Update objectness IoU sort (#3786)

Browse files
Files changed (1) hide show
  1. 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
- sort_id = torch.argsort(score_iou)
138
- b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[sort_id], gi[sort_id], score_iou[sort_id]
 
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