update o2o

yolo/models/yolov10/loss.py

@@ -70,7 +70,7 @@ class SetCriterion(object):
 
         return loss_dfl
 
-    def __call__(self, outputs, targets):        
+    def compute_loss(self, outputs, targets):        
         """
             outputs['pred_cls']: List(Tensor) [B, M, C]
             outputs['pred_reg']: List(Tensor) [B, M, 4*(reg_max+1)]
@@ -173,7 +173,22 @@ class SetCriterion(object):
         )
 
         return loss_dict
-    
+
+    def __call__(self, outputs, targets):
+        loss_o2o = self.compute_loss(outputs["outputs_o2o"], targets)
+        loss_o2m = self.compute_loss(outputs["outputs_o2m"], targets)
+
+        loss_dict = {}
+        for k in loss_o2o:
+            loss_dict[k+"_o2o"] = loss_o2o[k]
+
+        for k in loss_o2m:
+            loss_dict[k+"_o2m"] = loss_o2m[k]
+
+        loss_dict["losses"] = loss_o2o["losses"] + loss_o2m["losses"]
+
+        return loss_dict
+
 
 if __name__ == "__main__":
     pass

yolo/models/yolov10/yolov10.py

@@ -1,4 +1,5 @@
 # --------------- Torch components ---------------
+import copy
 import torch
 import torch.nn as nn
 
@@ -37,8 +38,11 @@ class Yolov10(nn.Module):
         self.fpn = Yolov10PaFPN(cfg, self.backbone.feat_dims)
 
         ## Head
-        self.head = Yolov10DetHead(cfg, self.fpn.out_dims)
-        self.pred = Yolov10DetPredLayer(cfg, self.head.cls_head_dim, self.head.reg_head_dim)
+        self.head_o2m = Yolov10DetHead(cfg, self.fpn.out_dims)
+        self.pred_o2m = Yolov10DetPredLayer(cfg, self.head.cls_head_dim, self.head.reg_head_dim)
+
+        self.head_o2o = copy.deepcopy(self.head_o2m)
+        self.pred_o2o = copy.deepcopy(self.pred_o2m)
 
     def post_process(self, cls_preds, box_preds):
         """
@@ -125,16 +129,15 @@ class Yolov10(nn.Module):
         # ---------------- PaFPN ----------------
         pyramid_feats = self.fpn(pyramid_feats)
 
-        # ---------------- Heads ----------------
-        cls_feats, reg_feats = self.head(pyramid_feats)
-
-        # ---------------- Preds ----------------
-        outputs = self.pred(cls_feats, reg_feats)
-        outputs['image_size'] = [x.shape[2], x.shape[3]]
+        # ---------------- Heads (one-to-one) ----------------
+        pyramid_feats_detach = [feat.detach() for feat in pyramid_feats]
+        cls_feats, reg_feats = self.head_o2o(pyramid_feats_detach)
+        outputs_o2o = self.pred_o2o(cls_feats, reg_feats)
+        outputs_o2o['image_size'] = [x.shape[2], x.shape[3]]
 
         if not self.training:
-            all_cls_preds = outputs['pred_cls']
-            all_box_preds = outputs['pred_box']
+            all_cls_preds = outputs_o2o['pred_cls']
+            all_box_preds = outputs_o2o['pred_box']
 
             # post process
             bboxes, scores, labels = self.post_process(all_cls_preds, all_box_preds)
@@ -143,5 +146,16 @@ class Yolov10(nn.Module):
                 "labels": labels,
                 "bboxes": bboxes
             }
-        
+        else:
+            # ---------------- Heads (one-to-many) ----------------
+            cls_feats, reg_feats = self.head_o2m(pyramid_feats)
+            outputs_o2m = self.pred_o2m(cls_feats, reg_feats)
+            outputs_o2m['image_size'] = [x.shape[2], x.shape[3]]
+
+            outputs = {
+                "outputs_o2o": outputs_o2o,
+                "outputs_o2m": outputs_o2m,
+            }
+            
+
         return outputs