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

2020-08-01 11:42:41 -07:00 · 2020-08-01 11:42:41 -07:00 · 23d2bde01f
parent bcd452c482
commit 23d2bde01f
1 changed files with 2 additions and 1 deletions
--- a/utils/utils.py
+++ b/utils/utils.py
@ -493,7 +493,8 @@ def compute_loss(p, targets, model):  # predictions, targets, model
    s = 3 / np  # output count scaling
    lbox *= h['giou'] * s
    lobj *= h['obj'] * s * (1.4 if np == 4 else 1.)
-    lcls *= h['cls'] * s
+    if model.nc > 1:
+        lcls *= h['cls'] * s
    bs = tobj.shape[0]  # batch size

    loss = lbox + lobj + lcls