updates

train.py

@@ -15,7 +15,7 @@ from utils.utils import *
 hyp = {'giou': 1.666,  # giou loss gain
        'xy': 4.062,  # xy loss gain
        'wh': 0.1845,  # wh loss gain
-       'cls': 1.0,  # cls loss gain
+       'cls': 42.6,  # cls loss gain
        'cls_pw': 3.34,  # cls BCELoss positive_weight
        'obj': 12.61,  # obj loss gain
        'obj_pw': 8.338,  # obj BCELoss positive_weight

utils/utils.py

@@ -281,7 +281,7 @@ def compute_loss(p, targets, model, giou_loss=True):  # predictions, targets, mo
     MSE = nn.MSELoss()
     BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']]))
     BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']]))
-    CE = nn.CrossEntropyLoss()  # (weight=model.class_weights)
+    # CE = nn.CrossEntropyLoss()  # (weight=model.class_weights)
 
     # Compute losses
     bs = p[0].shape[0]  # batch size
@@ -304,10 +304,10 @@ def compute_loss(p, targets, model, giou_loss=True):  # predictions, targets, mo
                 lxy += (k * h['xy']) * MSE(torch.sigmoid(pi[..., 0:2]), txy[i])  # xy loss
                 lwh += (k * h['wh']) * MSE(pi[..., 2:4], twh[i])  # wh yolo loss
 
-            # tclsm = torch.zeros_like(pi[..., 5:])
-            # tclsm[range(len(b)), tcls[i]] = 1.0
-            # lcls += (k * h['cls']) * BCEcls(pi[..., 5:], tclsm)  # cls loss (BCE)
-            lcls += (k * h['cls']) * CE(pi[..., 5:], tcls[i])  # cls loss (CE)
+            tclsm = torch.zeros_like(pi[..., 5:])
+            tclsm[range(len(b)), tcls[i]] = 1.0
+            lcls += (k * h['cls']) * BCEcls(pi[..., 5:], tclsm)  # cls loss (BCE)
+            # lcls += (k * h['cls']) * CE(pi[..., 5:], tcls[i])  # cls loss (CE)
 
             # Append targets to text file
             # with open('targets.txt', 'a') as file: