This commit is contained in:
Glenn Jocher 2019-06-12 19:40:21 +02:00
parent 1e8df4db23
commit 19d2232665
1 changed files with 1 additions and 3 deletions

View File

@ -274,6 +274,7 @@ def compute_loss(p, targets, model): # 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
lxy, lwh, lcls, lconf, lgiou = ft([0]), ft([0]), ft([0]), ft([0]), ft([0]) lxy, lwh, lcls, lconf, lgiou = ft([0]), ft([0]), ft([0]), ft([0]), ft([0])
txy, twh, tcls, tbox, indices, anchor_vec = build_targets(model, targets) txy, twh, tcls, tbox, indices, anchor_vec = build_targets(model, targets)
h = model.hyp # hyperparameters
# Define criteria # Define criteria
MSE = nn.MSELoss() MSE = nn.MSELoss()
@ -281,7 +282,6 @@ def compute_loss(p, targets, model): # predictions, targets, model
BCE = nn.BCEWithLogitsLoss() BCE = nn.BCEWithLogitsLoss()
# Compute losses # Compute losses
h = model.hyp # hyperparameters
bs = p[0].shape[0] # batch size bs = p[0].shape[0] # batch size
k = bs # loss gain k = bs # loss gain
for i, pi0 in enumerate(p): # layer i predictions, i for i, pi0 in enumerate(p): # layer i predictions, i
@ -303,8 +303,6 @@ def compute_loss(p, targets, model): # predictions, targets, model
lwh += (k * h['wh']) * MSE(pi[..., 2:4], twh[i]) # wh yolo loss lwh += (k * h['wh']) * MSE(pi[..., 2:4], twh[i]) # wh yolo loss
lcls += (k * h['cls']) * CE(pi[..., 5:], tcls[i]) # class_conf loss lcls += (k * h['cls']) * CE(pi[..., 5:], tcls[i]) # class_conf loss
# pos_weight = ft([gp[i] / min(gp) * 4.])
# BCE = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
lconf += (k * h['conf']) * BCE(pi0[..., 4], tconf) # obj_conf loss lconf += (k * h['conf']) * BCE(pi0[..., 4], tconf) # obj_conf loss
loss = lxy + lwh + lconf + lcls loss = lxy + lwh + lconf + lcls