assigner: cast to float32 as well in case you're training with bf16. Added focal loss

Author: me

src/assigner.cpp

@@ -266,9 +266,9 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> tal (
 
     // Store device for later then put everything on CPU
     auto device    = pred_boxes.device();
-    pred_boxes     = pred_boxes.to(torch::TensorOptions(torch::Device("cpu")));
-    pred_scores    = pred_scores.to(torch::TensorOptions(torch::Device("cpu")));
-    targets        = targets.to(torch::TensorOptions(torch::Device("cpu")));
+    pred_boxes     = pred_boxes.to(torch::kCPU, torch::kFloat);
+    pred_scores    = pred_scores.to(torch::kCPU, torch::kFloat);
+    targets        = targets.to(torch::kCPU, torch::kFloat);
 
     // Outputs 
     auto boxes      = torch::zeros({B, N, 4});

src/models.py

@@ -107,6 +107,9 @@ def exists(val):
 def default(val, d):
     return val if exists(val) else d
 
+def default_lazy(x, fn):
+    return x if exists(x) else fn()
+
 def Repeat(module, N):
     return nn.Sequential(*[deepcopy(module) for _ in range(N)])
 
@@ -886,6 +889,15 @@ def box2dist(box, sxy, strides):
     dist       = torch.cat([lt,rb], -1)
     return dist
 
+def focal_loss(logits, scores, soft_targets, gamma: float = 2.0, reduction: str = 'sum'):
+    p       = default_lazy(scores, lambda: logits.sigmoid())
+    bce     = F.binary_cross_entropy_with_logits(logits, soft_targets, reduction="none")
+    mod     = (p - soft_targets).abs().pow(gamma)
+    loss    = bce*mod
+    if   reduction == 'sum':  loss = loss.sum()
+    elif reduction == 'mean': loss = loss.mean()
+    return loss
+
 @torch.no_grad()
 def make_anchors(feats, strides): # anchor-free
     xys, strides2 = [], []
@@ -1028,7 +1040,8 @@ def forward_private(self, xs, cv2, cv3, targets=None):
             else:        loss_dfl = (F.l1_loss(ltrb[mask], box2dist(tboxes, sxy, strides)[mask], reduction='none') * weight.unsqueeze(-1)).sum() / tgt_scores_sum
 
             # Class loss (positive samples + negative)
-            loss_cls = F.binary_cross_entropy_with_logits(logits, tcls*tscores.unsqueeze(-1), reduction='sum') / tgt_scores_sum
+            # loss_cls = F.binary_cross_entropy_with_logits(logits, tcls*tscores.unsqueeze(-1), reduction='sum') / tgt_scores_sum
+            loss_cls = focal_loss(logits, probs, tcls*tscores.unsqueeze(-1), gamma=2.0, reduction='sum') / tgt_scores_sum
 
         return pred if not exists(targets) else (pred, {'iou': loss_iou, 'dfl': loss_dfl, 'cls': loss_cls})