This commit is contained in:
Glenn Jocher 2019-12-08 20:15:25 -08:00
parent 2300cb964a
commit 2391996474
1 changed files with 13 additions and 4 deletions

View File

@ -334,10 +334,11 @@ def compute_loss(p, targets, model): # predictions, targets, model
tcls, tbox, indices, anchor_vec = build_targets(model, targets) tcls, tbox, indices, anchor_vec = build_targets(model, targets)
h = model.hyp # hyperparameters h = model.hyp # hyperparameters
arc = model.arc # # (default, uCE, uBCE) detection architectures arc = model.arc # # (default, uCE, uBCE) detection architectures
red = 'mean' # Loss reduction (sum or mean)
# Define criteria # Define criteria
BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']])) BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']]), reduction=red)
BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']])) BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']]), reduction=red)
BCE = nn.BCEWithLogitsLoss() BCE = nn.BCEWithLogitsLoss()
CE = nn.CrossEntropyLoss() # weight=model.class_weights CE = nn.CrossEntropyLoss() # weight=model.class_weights
@ -346,13 +347,16 @@ def compute_loss(p, targets, model): # predictions, targets, model
BCEcls, BCEobj, BCE, CE = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g), FocalLoss(BCE, g), FocalLoss(CE, g) BCEcls, BCEobj, BCE, CE = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g), FocalLoss(BCE, g), FocalLoss(CE, g)
# Compute losses # Compute losses
np, ng = 0, 0 # number grid points, targets
for i, pi in enumerate(p): # layer index, layer predictions for i, pi in enumerate(p): # layer index, layer predictions
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
tobj = torch.zeros_like(pi[..., 0]) # target obj tobj = torch.zeros_like(pi[..., 0]) # target obj
np += tobj.numel()
# Compute losses # Compute losses
nb = len(b) nb = len(b)
if nb: # number of targets if nb: # number of targets
ng += nb
ps = pi[b, a, gj, gi] # prediction subset corresponding to targets ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
tobj[b, a, gj, gi] = 1.0 # obj tobj[b, a, gj, gi] = 1.0 # obj
# ps[:, 2:4] = torch.sigmoid(ps[:, 2:4]) # wh power loss (uncomment) # ps[:, 2:4] = torch.sigmoid(ps[:, 2:4]) # wh power loss (uncomment)
@ -360,8 +364,8 @@ def compute_loss(p, targets, model): # predictions, targets, model
# GIoU # GIoU
pxy = torch.sigmoid(ps[:, 0:2]) # pxy = pxy * s - (s - 1) / 2, s = 1.5 (scale_xy) pxy = torch.sigmoid(ps[:, 0:2]) # pxy = pxy * s - (s - 1) / 2, s = 1.5 (scale_xy)
pbox = torch.cat((pxy, torch.exp(ps[:, 2:4]).clamp(max=1E3) * anchor_vec[i]), 1) # predicted box pbox = torch.cat((pxy, torch.exp(ps[:, 2:4]).clamp(max=1E3) * anchor_vec[i]), 1) # predicted box
giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou computation giou = 1.0 - bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou computation
lbox += (1.0 - giou).mean() # giou loss lbox += giou.sum() if red == 'sum' else giou.mean() # giou loss
if 'default' in arc and model.nc > 1: # cls loss (only if multiple classes) if 'default' in arc and model.nc > 1: # cls loss (only if multiple classes)
t = torch.zeros_like(ps[:, 5:]) # targets t = torch.zeros_like(ps[:, 5:]) # targets
@ -396,6 +400,11 @@ def compute_loss(p, targets, model): # predictions, targets, model
lbox *= h['giou'] lbox *= h['giou']
lobj *= h['obj'] lobj *= h['obj']
lcls *= h['cls'] lcls *= h['cls']
if red == 'sum':
lbox *= 3 / ng
lobj *= 3 / np
lcls *= 3 / ng / model.nc
loss = lbox + lobj + lcls loss = lbox + lobj + lcls
return loss, torch.cat((lbox, lobj, lcls, loss)).detach() return loss, torch.cat((lbox, lobj, lcls, loss)).detach()