leaf variable --single-cls training bug fix (#593)

This commit is contained in:
Glenn Jocher committed 2020-08-01 11:42:41 -07:00
1 parent bcd452c482
commit 23d2bde01f
1 file changed
+1
+1
View File
@@ -493,6 +493,7 @@ def compute_loss(p, targets, model): # predictions, targets, model
s = 3 / np # output count scaling s = 3 / np # output count scaling
lbox *= h['giou'] * s lbox *= h['giou'] * s
lobj *= h['obj'] * s * (1.4 if np == 4 else 1.) lobj *= h['obj'] * s * (1.4 if np == 4 else 1.)
if model.nc > 1:
lcls *= h['cls'] * s lcls *= h['cls'] * s
bs = tobj.shape[0] # batch size bs = tobj.shape[0] # batch size