Yolo26 - inference done

Author: pfeatherstone

src/models.py

@@ -279,7 +279,7 @@ def forward(self, x):
         return x
 
 class SPPF(nn.Module):
-    def __init__(self, c1, c2, acts=[nn.Identity(), nn.SiLU(True)], shortcut=False):  # equivalent to SPP(k=(5, 9, 13))
+    def __init__(self, c1, c2, acts=[nn.SiLU(True), nn.SiLU(True)], shortcut=False):  # equivalent to SPP(k=(5, 9, 13))
         super().__init__()
         c_          = c1 // 2  # hidden channels
         self.add    = shortcut and c1==c2
@@ -567,7 +567,7 @@ def forward(self, x):
         return x4, x6, x10
 
 class BackboneV11(nn.Module):
-    def __init__(self, w, r, d, variant, sppf_shortcut=False):
+    def __init__(self, w, r, d, variant, sppf_shortcut=False, sppf_acts=[nn.SiLU(True), nn.SiLU(True)]):
         super().__init__()
         c3k = variant in "mlx"
         self.b0 = Conv(c1=3,            c2=int(64*w),    k=3, s=2)
@@ -579,7 +579,7 @@ def __init__(self, w, r, d, variant, sppf_shortcut=False):
         self.b6 = C3k2(c1=int(512*w),   c2=int(512*w),   n=round(2*d), e=0.50, c3k=True)
         self.b7 = Conv(c1=int(512*w),   c2=int(512*w*r), k=3, s=2)
         self.b8 = C3k2(c1=int(512*w*r), c2=int(512*w*r), n=round(2*d), e=0.50, c3k=True)
-        self.b9 = SPPF(c1=int(512*w*r), c2=int(512*w*r), shortcut=sppf_shortcut)
+        self.b9 = SPPF(c1=int(512*w*r), c2=int(512*w*r), shortcut=sppf_shortcut, acts=sppf_acts)
         self.b10 = PSA(int(512*w*r), n=round(2*d))
 
     def forward(self, x):
@@ -988,6 +988,7 @@ def __init__(self, nc=80, ch=(), dfl=True, separable=False, end2end=False):
         def spconv(c1, c2, k): return nn.Sequential(Conv(c1,c1,k,g=c1),Conv(c1,c2,1))
         conv = spconv if separable else Conv
         self.nc         = nc                        # number of classes
+        self.dfl        = dfl
         self.reg_max    = 16 if dfl else 1          # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
         self.no         = nc + self.reg_max * 4     # number of outputs per anchor
         self.strides    = [8, 16, 32]               # strides computed during build
@@ -1004,7 +1005,7 @@ def forward_private(self, xs, cv2, cv3, targets=None):
         feats       = [rearrange(torch.cat((c1(x), c2(x)), 1), 'b f h w -> b (h w) f') for x,c1,c2 in zip(xs, cv2, cv3)]
         dist, cls   = torch.cat(feats, 1).split((4 * self.reg_max, self.nc), -1)
         dist        = rearrange(dist, 'b n (k r) -> b n k r', k=4)
-        ltrb        = torch.einsum('bnkr, r -> bnk', dist.softmax(-1), self.r)
+        ltrb        = torch.einsum('bnkr, r -> bnk', dist.softmax(-1), self.r) if self.dfl else dist.squeeze(-1)
         box         = dist2box(ltrb, sxy, strides)
         pred        = torch.cat((box, cls.sigmoid()), -1)
 
@@ -1177,7 +1178,7 @@ def __init__(self, variant, num_classes):
 class Yolov26(YoloBase):
     def __init__(self, variant, num_classes):
         d, w, r = get_variant_multiplesV26(variant)
-        super().__init__(BackboneV11(w, r, d, variant, sppf_shortcut=True),
+        super().__init__(BackboneV11(w, r, d, variant, sppf_shortcut=True, sppf_acts=[nn.Identity(), nn.SiLU(True)]),
                          HeadV11(w, r, d, variant, is26=True),
                          Detect(num_classes, ch=(int(256*w), int(512*w), int(512*w*r)), separable=True, dfl=False, end2end=True),
                          variant)

src/test.py

@@ -141,6 +141,9 @@ def load_from_ultralytics(net: Union[Yolov5, Yolov8, Yolov10, Yolov11]):
         for module in net2.modules():
             if isinstance(module, AAttn):
                 fuse_bias_v12(module.pe.conv, module.pe.bn)
+    elif isinstance(net, Yolov26):
+        net2  = YOLO('yolo26{}.pt'.format(net.v)).model.eval()
+        l0,l1 = 11,23
 
     assert (nP1 := count_parameters(net)) == (nP2 := count_parameters(net2)), f'wrong number of parameters net {nP1} vs ultralytics {nP2}'
     copy_params(net.net, net2.model[0:l0])
@@ -187,8 +190,6 @@ def params(n):
         assert p1.shape == p2.shape, f"bad shape: {k} {p2.shape} {p1.shape}"
         p1.data.copy_(p2.data)
 
-    init_batchnorms(net)
-
 
 def load_from_yolov7_official(net: Yolov7, weights_pt: str):
     def params1():
@@ -206,8 +207,6 @@ def params2():
     for p1, p2 in zip(params1(), params2(), strict=True):
         p1.data.copy_(p2.data)
 
-    init_batchnorms(net)
-
     # Handle special case in SPPCSPC where Yolov6 and Yolov7 disagree on the order of the final torch.cat()
     for module in net.modules():
         if isinstance(module, SPPCSPC):
@@ -230,6 +229,7 @@ def get_model(model: str, variant: str = ''):
         case 'yolov10':     net = Yolov10(variant, 80).eval()
         case 'yolov11':     net = Yolov11(variant, 80).eval()
         case 'yolov12':     net = Yolov12(variant, 80).eval()
+        case 'yolov26':     net = Yolov26(variant, 80).eval()
 
     print(f"{model}{variant} has {count_parameters(net)} parameters")
 
@@ -241,7 +241,7 @@ def get_model(model: str, variant: str = ''):
         download_if_not_exist(model, filepath)
         load_from_darknet(net, filepath)
 
-    if model in ['yolov5', 'yolov8', 'yolov10', 'yolov11', 'yolov12']:
+    if model in ['yolov5', 'yolov8', 'yolov10', 'yolov11', 'yolov12', 'yolov26']:
         load_from_ultralytics(net)
         has_obj = False
 
@@ -293,7 +293,6 @@ def export(model: str, variant: str = '', onnx_path:str = '/tmp/model.onnx'):
     torch.testing.assert_close(preds1[...,4:], torch.from_numpy(preds2[...,4:]))                        # scores
     print(bcolors.OKGREEN, "Checking with onnxruntime... Done", bcolors.ENDC)
 
-
 test('yolov3')
 test('yolov3-spp')
 test('yolov3-tiny')
@@ -329,6 +328,11 @@ def export(model: str, variant: str = '', onnx_path:str = '/tmp/model.onnx'):
 test('yolov12', 'm')
 test('yolov12', 'l')
 test('yolov12', 'x')
+test('yolov26', 'n')
+test('yolov26', 's')
+test('yolov26', 'm')
+test('yolov26', 'l')
+test('yolov26', 'x')
 
 # export('yolov3')
 # export('yolov3-spp')
@@ -364,4 +368,4 @@ def export(model: str, variant: str = '', onnx_path:str = '/tmp/model.onnx'):
 # export('yolov12', 's')
 # export('yolov12', 'm')
 # export('yolov12', 'l')
-# export('yolov12', 'x')
+# export('yolov12', 'x')