Seg loss afects whole network ultralytics#9

manole-alexandru · Mar 26, 2023 · bf0183b · bf0183b
1 parent 25b23f6
commit bf0183b
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/train.py b/train.py
@@ -513,7 +513,7 @@ def parse_opt(known=False):
 
 
 def main(opt, callbacks=Callbacks()):
-    print('\n---------- VERSION:', '#0008', '----------\n')
+    print('\n---------- VERSION:', '#0009', '----------\n')
     # Checks
     if RANK in {-1, 0}:
         print_args(vars(opt))

diff --git a/utils/loss.py b/utils/loss.py
@@ -89,7 +89,7 @@ def forward(self, pred, true):
 
 def weighted_bce(y_pred, y_true, BETA=2):
     weights = (y_true * (BETA - 1)) + 1
-    bce = nn.BCELoss(reduction='none')(y_pred, y_true)
+    bce = nn.BCEWithLogitsLoss(reduction='none')(y_pred, y_true)
     wbce = torch.mean(bce * weights)
     return wbce
 

diff --git a/utils/torch_utils.py b/utils/torch_utils.py
@@ -355,8 +355,8 @@ def smart_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5):
                     g_det[1].append(p)
                 else:
                     g_det[0].append(p)  # weight (with decay)
-        else:
-        # if m_index < 287 or m_index > 290: # SEGMENTATION Optimizer
+        # else:
+        if m_index < 287 or m_index > 290: # SEGMENTATION Optimizer
             for p_index, (p_name, p) in enumerate(v.named_parameters(recurse=0)):
                 if p_name == 'bias':  # bias (no decay)
                     g_seg[2].append(p)