diff --git a/models.py b/models.py index 5df79dc7..fb87a649 100755 --- a/models.py +++ b/models.py @@ -82,7 +82,7 @@ def create_modules(module_defs, img_size, arc): elif arc == 'uCE': # unified CE (1 background + 80 classes) b = [7, 0] # obj, cls elif arc == 'uBCE': # unified BCE (80 classes) - b = [0, -4] # obj, cls + b = [0, -3.5] # obj, cls bias = module_list[-1][0].bias.view(len(mask), -1) # 255 to 3x85 bias[:, 4] += b[0] # obj diff --git a/utils/utils.py b/utils/utils.py index ef175c10..d94d91ab 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -322,8 +322,8 @@ def compute_loss(p, targets, model): # predictions, targets, model # Define criteria BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']])) BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']])) - FBCE = FocalLoss(nn.BCEWithLogitsLoss()) - FCE = FocalLoss(nn.CrossEntropyLoss()) # weight=model.class_weights + FBCE = nn.BCEWithLogitsLoss() + FCE = nn.CrossEntropyLoss() # weight=model.class_weights # Compute losses for i, pi in enumerate(p): # layer index, layer predictions