update

yolo/models/yolo11/modules.py

@@ -150,7 +150,7 @@ class C2PSA(nn.Module):
 
 
 # ----------------- YOLO11 components -----------------
-class YoloStage(nn.Module):
+class C3k2fBlock(nn.Module):
     def __init__(self, in_dim, out_dim, num_blocks=1, use_c3k=True, expansion=0.5, shortcut=True):
         super().__init__()
         inter_dim = int(out_dim * expansion)  # hidden channels

yolo/models/yolo11/yolo11_backbone.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .modules import ConvModule, YoloStage, SPPF, C2PSA
+    from .modules import ConvModule, C3k2fBlock, SPPF, C2PSA
 except:
-    from  modules import ConvModule, YoloStage, SPPF, C2PSA
+    from  modules import ConvModule, C3k2fBlock, SPPF, C2PSA
 
 
 # ---------------------------- YOLO11 Backbone ----------------------------
@@ -21,46 +21,46 @@ class Yolo11Backbone(nn.Module):
         # P2/4
         self.layer_2 = nn.Sequential(
             ConvModule(int(64 * cfg.width), int(128 * cfg.width), kernel_size=3, stride=2),
-            YoloStage(in_dim     = int(128 * cfg.width),
-                      out_dim    = int(256 * cfg.width),
-                      num_blocks = round(2*cfg.depth),
-                      shortcut   = True,
-                      expansion  = 0.25,
-                      use_c3k    = False if self.model_scale in "ns" else True,
-                      )
+            C3k2fBlock(in_dim     = int(128 * cfg.width),
+                       out_dim    = int(256 * cfg.width),
+                       num_blocks = round(2*cfg.depth),
+                       shortcut   = True,
+                       expansion  = 0.25,
+                       use_c3k    = False if self.model_scale in "ns" else True,
+                       )
         )
         # P3/8
         self.layer_3 = nn.Sequential(
             ConvModule(int(256 * cfg.width), int(256 * cfg.width), kernel_size=3, stride=2),
-            YoloStage(in_dim     = int(256 * cfg.width),
-                      out_dim    = int(512 * cfg.width),
-                      num_blocks = round(2*cfg.depth),
-                      shortcut   = True,
-                      expansion  = 0.25,
-                      use_c3k    = False if self.model_scale in "ns" else True,
-                      )
+            C3k2fBlock(in_dim     = int(256 * cfg.width),
+                       out_dim    = int(512 * cfg.width),
+                       num_blocks = round(2*cfg.depth),
+                       shortcut   = True,
+                       expansion  = 0.25,
+                       use_c3k    = False if self.model_scale in "ns" else True,
+                       )
         )
         # P4/16
         self.layer_4 = nn.Sequential(
             ConvModule(int(512 * cfg.width), int(512 * cfg.width), kernel_size=3, stride=2),
-            YoloStage(in_dim     = int(512 * cfg.width),
-                      out_dim    = int(512 * cfg.width),
-                      num_blocks = round(2*cfg.depth),
-                      shortcut   = True,
-                      expansion  = 0.5,
-                      use_c3k    = True,
-                      )
+            C3k2fBlock(in_dim     = int(512 * cfg.width),
+                       out_dim    = int(512 * cfg.width),
+                       num_blocks = round(2*cfg.depth),
+                       shortcut   = True,
+                       expansion  = 0.5,
+                       use_c3k    = True,
+                       )
         )
         # P5/32
         self.layer_5 = nn.Sequential(
             ConvModule(int(512 * cfg.width), int(512 * cfg.width * cfg.ratio), kernel_size=3, stride=2),
-            YoloStage(in_dim     = int(512 * cfg.width * cfg.ratio),
-                      out_dim    = int(512 * cfg.width * cfg.ratio),
-                      num_blocks = round(2*cfg.depth),
-                      shortcut   = True,
-                      expansion  = 0.5,
-                      use_c3k    = True,
-                      )
+            C3k2fBlock(in_dim     = int(512 * cfg.width * cfg.ratio),
+                       out_dim    = int(512 * cfg.width * cfg.ratio),
+                       num_blocks = round(2*cfg.depth),
+                       shortcut   = True,
+                       expansion  = 0.5,
+                       use_c3k    = True,
+                       )
         )
         # Extra module (no pretrained weight)
         self.layer_6 = SPPF(in_dim  = int(512 * cfg.width * cfg.ratio),

yolo/models/yolo11/yolo11_pafpn.py

@@ -4,9 +4,9 @@ import torch.nn.functional as F
 from typing import List
 
 try:
-    from .modules import ConvModule, YoloStage
+    from .modules import ConvModule, C3k2fBlock
 except:
-    from  modules import ConvModule, YoloStage
+    from  modules import ConvModule, C3k2fBlock
 
 
 class Yolo11PaFPN(nn.Module):
@@ -19,40 +19,40 @@ class Yolo11PaFPN(nn.Module):
 
         # ----------------------------- Yolo11's Top-down FPN -----------------------------
         ## P5 -> P4
-        self.top_down_layer_1 = YoloStage(in_dim     = self.in_dims[0] + self.in_dims[1],
-                                          out_dim    = round(512*cfg.width),
-                                          num_blocks = round(2 * cfg.depth),
-                                          shortcut   = True,
-                                          expansion  = 0.5,
-                                          use_c3k    = False if self.model_scale in "ns" else True,
-                                          )
-        ## P4 -> P3
-        self.top_down_layer_2 = YoloStage(in_dim     = self.in_dims[2] + round(512*cfg.width),
-                                          out_dim    = round(256*cfg.width),
-                                          num_blocks = round(2 * cfg.depth),
-                                          shortcut   = True,
-                                          expansion  = 0.5,
-                                          use_c3k    = False if self.model_scale in "ns" else True,
-                                          )
-        # ----------------------------- Yolo11's Bottom-up PAN -----------------------------
-        ## P3 -> P4
-        self.dowmsample_layer_1 = ConvModule(round(256*cfg.width), round(256*cfg.width), kernel_size=3, stride=2)
-        self.bottom_up_layer_1 = YoloStage(in_dim     = round(256*cfg.width) + round(512*cfg.width),
+        self.top_down_layer_1 = C3k2fBlock(in_dim     = self.in_dims[0] + self.in_dims[1],
                                            out_dim    = round(512*cfg.width),
                                            num_blocks = round(2 * cfg.depth),
                                            shortcut   = True,
                                            expansion  = 0.5,
                                            use_c3k    = False if self.model_scale in "ns" else True,
                                            )
-        ## P4 -> P5
-        self.dowmsample_layer_2 = ConvModule(round(512*cfg.width), round(512*cfg.width), kernel_size=3, stride=2)
-        self.bottom_up_layer_2 = YoloStage(in_dim     = round(512*cfg.width) + self.in_dims[0],
-                                           out_dim    = round(512*cfg.width*cfg.ratio),
+        ## P4 -> P3
+        self.top_down_layer_2 = C3k2fBlock(in_dim     = self.in_dims[2] + round(512*cfg.width),
+                                           out_dim    = round(256*cfg.width),
                                            num_blocks = round(2 * cfg.depth),
                                            shortcut   = True,
                                            expansion  = 0.5,
-                                           use_c3k    = True,
+                                           use_c3k    = False if self.model_scale in "ns" else True,
                                            )
+        # ----------------------------- Yolo11's Bottom-up PAN -----------------------------
+        ## P3 -> P4
+        self.dowmsample_layer_1 = ConvModule(round(256*cfg.width), round(256*cfg.width), kernel_size=3, stride=2)
+        self.bottom_up_layer_1 = C3k2fBlock(in_dim     = round(256*cfg.width) + round(512*cfg.width),
+                                            out_dim    = round(512*cfg.width),
+                                            num_blocks = round(2 * cfg.depth),
+                                            shortcut   = True,
+                                            expansion  = 0.5,
+                                            use_c3k    = False if self.model_scale in "ns" else True,
+                                            )
+        ## P4 -> P5
+        self.dowmsample_layer_2 = ConvModule(round(512*cfg.width), round(512*cfg.width), kernel_size=3, stride=2)
+        self.bottom_up_layer_2 = C3k2fBlock(in_dim     = round(512*cfg.width) + self.in_dims[0],
+                                            out_dim    = round(512*cfg.width*cfg.ratio),
+                                            num_blocks = round(2 * cfg.depth),
+                                            shortcut   = True,
+                                            expansion  = 0.5,
+                                            use_c3k    = True,
+                                            )
 
         self.init_weights()