align loss to darknet
This commit is contained in:
parent
7416c1842a
commit
ff630b1960
|
@ -139,7 +139,7 @@ class YOLOLayer(nn.Module):
|
|||
if targets is not None:
|
||||
MSELoss = nn.MSELoss(size_average=True)
|
||||
BCEWithLogitsLoss = nn.BCEWithLogitsLoss(size_average=True)
|
||||
# CrossEntropyLoss = nn.CrossEntropyLoss()
|
||||
CrossEntropyLoss = nn.CrossEntropyLoss()
|
||||
|
||||
if requestPrecision:
|
||||
gx = self.grid_x[:, :, :nG, :nG]
|
||||
|
@ -170,8 +170,8 @@ class YOLOLayer(nn.Module):
|
|||
# lconf = k * BCEWithLogitsLoss(pred_conf[mask], mask[mask].float())
|
||||
lconf = k * BCEWithLogitsLoss(pred_conf, mask.float())
|
||||
|
||||
# lcls = k * CrossEntropyLoss(pred_cls[mask], torch.argmax(tcls, 1))
|
||||
lcls = k * BCEWithLogitsLoss(pred_cls[mask], tcls.float())
|
||||
lcls = k * CrossEntropyLoss(pred_cls[mask], torch.argmax(tcls, 1))
|
||||
# lcls = k * BCEWithLogitsLoss(pred_cls[mask], tcls.float())
|
||||
else:
|
||||
lx, ly, lw, lh, lcls, lconf = FT([0]), FT([0]), FT([0]), FT([0]), FT([0]), FT([0])
|
||||
|
||||
|
|
Loading…
Reference in New Issue