This commit is contained in:
Glenn Jocher 2020-01-17 23:30:17 -08:00
parent 3bac3c63b1
commit 43956d6305
2 changed files with 4 additions and 4 deletions

View File

@ -221,7 +221,7 @@ def train():
# Prebias # Prebias
if prebias: if prebias:
if epoch < 3: # prebias if epoch < 1: # prebias
ps = 0.1, 0.9 # prebias settings (lr=0.1, momentum=0.9) ps = 0.1, 0.9 # prebias settings (lr=0.1, momentum=0.9)
else: # normal training else: # normal training
ps = hyp['lr0'], hyp['momentum'] # normal training settings ps = hyp['lr0'], hyp['momentum'] # normal training settings
@ -278,7 +278,7 @@ def train():
pred = model(imgs) pred = model(imgs)
# Compute loss # Compute loss
loss, loss_items = compute_loss(pred, targets, model) loss, loss_items = compute_loss(pred, targets, model, not prebias)
if not torch.isfinite(loss): if not torch.isfinite(loss):
print('WARNING: non-finite loss, ending training ', loss_items) print('WARNING: non-finite loss, ending training ', loss_items)
return results return results

View File

@ -361,7 +361,7 @@ class FocalLoss(nn.Module):
return loss return loss
def compute_loss(p, targets, model): # predictions, targets, model def compute_loss(p, targets, model, giou_flag=True): # predictions, targets, model
ft = torch.cuda.FloatTensor if p[0].is_cuda else torch.Tensor ft = torch.cuda.FloatTensor if p[0].is_cuda else torch.Tensor
lcls, lbox, lobj = ft([0]), ft([0]), ft([0]) lcls, lbox, lobj = ft([0]), ft([0]), ft([0])
tcls, tbox, indices, anchor_vec = build_targets(model, targets) tcls, tbox, indices, anchor_vec = build_targets(model, targets)
@ -399,7 +399,7 @@ def compute_loss(p, targets, model): # predictions, targets, model
pbox = torch.cat((pxy, pwh), 1) # predicted box pbox = torch.cat((pxy, pwh), 1) # predicted box
giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou computation 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 lbox += (1.0 - giou).sum() if red == 'sum' else (1.0 - giou).mean() # giou loss
tobj[b, a, gj, gi] = giou.detach().type(tobj.dtype) tobj[b, a, gj, gi] = giou.detach().type(tobj.dtype) if giou_flag else 1.0
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