update

yolo/models/yolov8/matcher.py

@@ -83,6 +83,7 @@ class TaskAlignedAssigner(nn.Module):
         overlaps[mask_in_gts] = bbox_iou(gt_boxes, pd_boxes, xywh=False, CIoU=True).squeeze(-1).clamp_(0)
 
         align_metric = bbox_scores.pow(self.alpha) * overlaps.pow(self.beta)
+        
         return align_metric, overlaps
 
     def select_topk_candidates(self, metrics, largest=True):

yolo/models/yolov9/gelan_neck.py

@@ -34,4 +34,33 @@ class SPPElan(nn.Module):
         y.extend(self.pool_layer(y[-1]) for _ in range(3))
         
         return self.conv_layer_2(torch.cat(y, 1))
-    
+
+
+if __name__=='__main__':
+    from thop import profile
+
+    class BaseConfig(object):
+        def __init__(self) -> None:
+            self.spp_inter_dim = 512
+            self.spp_out_dim = 512
+
+    # 定义模型配置文件
+    cfg = BaseConfig()
+
+    # Build a neck
+    in_dim  = 512
+    model = SPPElan(cfg, in_dim)
+
+    # Randomly generate a input data
+    x = torch.randn(2, in_dim, 20, 20)
+
+    # Inference
+    output = model(x)
+    print(' - the shape of input :  ', x.shape)
+    print(' - the shape of output : ', output.shape)
+
+    x = torch.randn(1, in_dim, 20, 20)
+    flops, params = profile(model, inputs=(x, ), verbose=False)
+    print('============== FLOPs & Params ================')
+    print(' - FLOPs  : {:.2f} G'.format(flops / 1e9 * 2))
+    print(' - Params : {:.2f} M'.format(params / 1e6))