align loss to darknet
This commit is contained in:
parent
313a3f6b0c
commit
292af1f2f4
|
@ -137,8 +137,8 @@ class YOLOLayer(nn.Module):
|
||||||
|
|
||||||
# Training
|
# Training
|
||||||
if targets is not None:
|
if targets is not None:
|
||||||
MSELoss = nn.MSELoss()
|
MSELoss = nn.MSELoss(size_average=False)
|
||||||
BCEWithLogitsLoss = nn.BCEWithLogitsLoss()
|
BCEWithLogitsLoss = nn.BCEWithLogitsLoss(size_average=False)
|
||||||
# CrossEntropyLoss = nn.CrossEntropyLoss()
|
# CrossEntropyLoss = nn.CrossEntropyLoss()
|
||||||
|
|
||||||
if requestPrecision:
|
if requestPrecision:
|
||||||
|
@ -160,7 +160,7 @@ class YOLOLayer(nn.Module):
|
||||||
nT = sum([len(x) for x in targets]) # number of targets
|
nT = sum([len(x) for x in targets]) # number of targets
|
||||||
nM = mask.sum().float() # number of anchors (assigned to targets)
|
nM = mask.sum().float() # number of anchors (assigned to targets)
|
||||||
nB = len(targets) # batch size
|
nB = len(targets) # batch size
|
||||||
k = nM / nB
|
k = 1 / nB
|
||||||
if nM > 0:
|
if nM > 0:
|
||||||
lx = k * MSELoss(x[mask], tx[mask])
|
lx = k * MSELoss(x[mask], tx[mask])
|
||||||
ly = k * MSELoss(y[mask], ty[mask])
|
ly = k * MSELoss(y[mask], ty[mask])
|
||||||
|
|
Loading…
Reference in New Issue