updates
This commit is contained in:
parent
1e8df4db23
commit
19d2232665
|
@ -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
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue