This commit is contained in:
Glenn Jocher
2020-03-04 13:20:08 -08:00
parent 6ab753a9e7
commit 981b452b1d
2 changed files with 4 additions and 3 deletions
+2 -2
View File
@@ -363,7 +363,7 @@ class FocalLoss(nn.Module):
return loss
def compute_loss(p, targets, model, giou_flag=True): # predictions, targets, model
def compute_loss(p, targets, model): # predictions, targets, model
ft = torch.cuda.FloatTensor if p[0].is_cuda else torch.Tensor
lcls, lbox, lobj = ft([0]), ft([0]), ft([0])
tcls, tbox, indices, anchor_vec = build_targets(model, targets)
@@ -401,7 +401,7 @@ def compute_loss(p, targets, model, giou_flag=True): # predictions, targets, mo
pbox = torch.cat((pxy, pwh), 1) # predicted box
giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou computation
lbox += (1.0 - giou).sum() if red == 'sum' else (1.0 - giou).mean() # giou loss
tobj[b, a, gj, gi] = giou.detach().clamp(0).type(tobj.dtype) if giou_flag else 1.0
tobj[b, a, gj, gi] = (1.0 - h['gr']) + h['gr'] * giou.detach().clamp(0).type(tobj.dtype) # giou ratio
if 'default' in arc and model.nc > 1: # cls loss (only if multiple classes)
t = torch.zeros_like(ps[:, 5:]) # targets