fix fp16 siou bug

bubbliiiing · bubbliiiing · commit d922f01030b9 · 2022-07-06T10:10:41.000+08:00
diff --git a/nets/yolo_training.py b/nets/yolo_training.py
@@ -125,7 +125,7 @@ def box_iou(self, b1, b2):
             #----------------------------------------------------#
             #   计算中心的距离
             #----------------------------------------------------#
-            sigma       = torch.pow(center_wh[..., 0] ** 2 + center_wh[..., 1] ** 2, 0.5)
+            sigma       = torch.pow(torch.sum(torch.pow(center_wh, 2), axis=-1), 0.5)
             
             #----------------------------------------------------#
             #   求h和w方向上的sin比值
@@ -264,8 +264,9 @@ def forward(self, l, input, targets=None):
             #   loss_cls 分类损失
             #---------------------------------------------------------------#
             iou         = self.box_iou(pred_boxes, y_true[..., :4]).type_as(x)
-            # loss_loc    = torch.mean((1 - iou)[obj_mask] * box_loss_scale[obj_mask])
+            obj_mask    = obj_mask & torch.logical_not(torch.isnan(iou))
             loss_loc    = torch.mean((1 - iou)[obj_mask])
+            # loss_loc    = torch.mean((1 - iou)[obj_mask] * box_loss_scale[obj_mask])
             
             loss_cls    = torch.mean(self.BCELoss(pred_cls[obj_mask], y_true[..., 5:][obj_mask]))
             loss        += loss_loc * self.box_ratio + loss_cls * self.cls_ratio
diff --git a/utils/utils_fit.py b/utils/utils_fit.py
@@ -60,7 +60,9 @@ def fit_one_epoch(model_train, model, yolo_loss, loss_history, eval_callback, op
                 #   计算损失
                 #----------------------#
                 for l in range(len(outputs)):
-                    loss_item = yolo_loss(l, outputs[l], targets)
+                    with torch.cuda.amp.autocast(enabled=False):
+                        predication = outputs[l].float()
+                    loss_item = yolo_loss(l, predication, targets)
                     loss_value_all  += loss_item
                 loss_value = loss_value_all
 
