YOLO26-style architecture

Important clarification first

YOLO26 is primarily a CNN architecture, not a pure Transformer detector (unlike DETR / RT-DETR).

It uses:

• C3k2 blocks (evolved from C2f / CSP)
• SPPF + residual shortcut
• C2PSA (Position-Sensitive Attention) — the main “attention / transformer-like” component
• Dual detection head (One-to-One NMS-free + One-to-Many)
• No DFL (Distribution Focal Loss removed)
• End-to-end NMS-free inference by default

Below is a complete, self-contained, educational PyTorch implementation of the YOLO26-style architecture so you can understand every block clearly.

pythonCopy
1"""
2Educational YOLO26-style Architecture (Ultralytics YOLO26 inspired)
3===================================================================
4Key features implemented for learning:
5  - C3k2 backbone/neck blocks
6  - SPPF with residual shortcut
7  - C2PSA (Position-Sensitive Attention)  ← the "transformer-like" part
8  - Dual-head: One-to-One (NMS-free) + One-to-Many
9  - DFL-free box regression (direct 4 values)
10  - Anchor-free, multi-scale detection (P3, P4, P5)
11
12This is a clean teaching version — not a 1:1 copy of Ultralytics source.
13"""
14
15import math
16import torch
17import torch.nn as nn
18import torch.nn.functional as F
19from typing import List, Tuple, Optional
20
21
22# =====================================================================
23# 1. Basic Building Blocks
24# =====================================================================
25
26def autopad(k, p=None):
27    """Auto padding so that 'same' convolution keeps spatial size."""
28    if p is None:
29        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
30    return p
31
32
33class Conv(nn.Module):
34    """Standard Conv-BN-SiLU."""
35    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
36        super().__init__()
37        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
38        self.bn = nn.BatchNorm2d(c2)
39        self.act = nn.SiLU() if act else nn.Identity()
40
41    def forward(self, x):
42        return self.act(self.bn(self.conv(x)))
43
44
45class Bottleneck(nn.Module):
46    """Standard bottleneck used inside C3k2."""
47    def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
48        super().__init__()
49        c_ = int(c2 * e)
50        self.cv1 = Conv(c1, c_, k[0], 1)
51        self.cv2 = Conv(c_, c2, k[1], 1, g=g)
52        self.add = shortcut and c1 == c2
53
54    def forward(self, x):
55        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
56
57
58# =====================================================================
59# 2. C3k2  (core block of YOLO11 / YOLO26)
60# =====================================================================
61
62class C3k(nn.Module):
63    """C3 variant with configurable kernel size (used when c3k=True)."""
64    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, k=3):
65        super().__init__()
66        c_ = int(c2 * e)
67        self.cv1 = Conv(c1, c_, 1, 1)
68        self.cv2 = Conv(c1, c_, 1, 1)
69        self.cv3 = Conv(2 * c_, c2, 1)
70        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=(k, k), e=1.0) for _ in range(n)))
71
72    def forward(self, x):
73        return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
74
75
76class C3k2(nn.Module):
77    """
78    C3k2 = improved C2f used in YOLO11 & YOLO26.
79    When c3k=False → plain Bottleneck
80    When c3k=True  → C3k block
81    """
82    def __init__(self, c1, c2, n=1, c3k=False, e=0.5, g=1, shortcut=True):
83        super().__init__()
84        self.c = int(c2 * e)          # hidden channels
85        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
86        self.cv2 = Conv((2 + n) * self.c, c2, 1)
87        self.m = nn.ModuleList(
88            C3k(self.c, self.c, 2, shortcut, g) if c3k
89            else Bottleneck(self.c, self.c, shortcut, g)
90            for _ in range(n)
91        )
92
93    def forward(self, x):
94        y = list(self.cv1(x).chunk(2, 1))
95        y.extend(m(y[-1]) for m in self.m)
96        return self.cv2(torch.cat(y, 1))
97
98
99# =====================================================================
100# 3. SPPF + Residual (YOLO26 improvement)
101# =====================================================================
102
103class SPPF(nn.Module):
104    """Spatial Pyramid Pooling - Fast + residual shortcut (YOLO26 style)."""
105    def __init__(self, c1, c2, k=5):
106        super().__init__()
107        c_ = c1 // 2
108        self.cv1 = Conv(c1, c_, 1, 1)
109        self.cv2 = Conv(c_ * 4, c2, 1, 1)
110        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
111        # residual projection if channels differ
112        self.shortcut = Conv(c1, c2, 1, 1) if c1 != c2 else nn.Identity()
113
114    def forward(self, x):
115        residual = self.shortcut(x)
116        x = self.cv1(x)
117        y1 = self.m(x)
118        y2 = self.m(y1)
119        y3 = self.m(y2)
120        out = self.cv2(torch.cat([x, y1, y2, y3], 1))
121        return out + residual          # residual connection (YOLO26)
122
123
124# =====================================================================
125# 4. C2PSA  – Position-Sensitive Attention (the "transformer" part)
126# =====================================================================
127
128class PSABlock(nn.Module):
129    """Single Position-Sensitive Attention block."""
130    def __init__(self, c, num_heads=4, attn_ratio=0.5):
131        super().__init__()
132        self.num_heads = num_heads
133        self.head_dim = c // num_heads
134        self.qkv = Conv(c, c * 3, 1, act=False)
135        self.proj = Conv(c, c, 1, act=False)
136        self.pe = Conv(c, c, 3, 1, g=c, act=False)   # position encoding (depthwise)
137
138    def forward(self, x):
139        B, C, H, W = x.shape
140        qkv = self.qkv(x).view(B, 3, self.num_heads, self.head_dim, H * W)
141        q, k, v = qkv.unbind(1)                       # each: (B, heads, head_dim, HW)
142
143        # scaled dot-product attention
144        attn = (q.transpose(-2, -1) @ k) * (self.head_dim ** -0.5)
145        attn = attn.softmax(dim=-1)
146        out = (v @ attn.transpose(-2, -1)).view(B, C, H, W)
147
148        # add local position encoding
149        out = out + self.pe(v.reshape(B, C, H, W))
150        return self.proj(out)
151
152
153class C2PSA(nn.Module):
154    """
155    C2PSA – CSP + Position Sensitive Attention
156    This is the main attention / transformer-style module in YOLO11/26.
157    """
158    def __init__(self, c1, c2, n=1, e=0.5):
159        super().__init__()
160        assert c1 == c2
161        self.c = int(c1 * e)
162        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
163        self.cv2 = Conv(2 * self.c, c1, 1)
164        self.m = nn.Sequential(*(PSABlock(self.c) for _ in range(n)))
165
166    def forward(self, x):
167        a, b = self.cv1(x).chunk(2, 1)
168        b = self.m(b)
169        return self.cv2(torch.cat((a, b), 1))
170
171
172# =====================================================================
173# 5. Detection Head  (DFL-free + Dual Head)
174# =====================================================================
175
176class Detect(nn.Module):
177    """
178    YOLO26-style dual detection head.
179    - One-to-Many  : classic dense predictions (needs NMS)
180    - One-to-One   : end-to-end NMS-free (max 300 boxes)
181    Box regression is DFL-free (direct 4 values: xywh or ltrb).
182    """
183    def __init__(self, nc=80, ch=()):
184        super().__init__()
185        self.nc = nc                  # number of classes
186        self.nl = len(ch)             # number of detection layers (P3, P4, P5)
187        self.reg_max = 1              # DFL removed → reg_max=1 (direct 4 values)
188        self.no = nc + 4              # outputs per anchor: 4 box + nc cls
189
190        # shared stem
191        c2, c3 = max(ch[0] // 4, 16), max(ch[0], min(nc * 2, 128))
192
193        self.cv2 = nn.ModuleList(     # box branch
194            nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4, 1))
195            for x in ch
196        )
197        self.cv3 = nn.ModuleList(     # class branch
198            nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, nc, 1))
199            for x in ch
200        )
201
202        # One-to-One head (for NMS-free inference) – copy of the branches
203        self.one2one_cv2 = nn.ModuleList(
204            nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4, 1))
205            for x in ch
206        )
207        self.one2one_cv3 = nn.ModuleList(
208            nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, nc, 1))
209            for x in ch
210        )
211
212        self.stride = torch.zeros(self.nl)
213        self.export = False
214        self.end2end = True           # default YOLO26 behaviour
215
216    def forward(self, x: List[torch.Tensor]):
217        """
218        x = list of feature maps [P3, P4, P5]
219        Returns:
220          - training : (one2many, one2one)
221          - inference (end2end=True) : one2one predictions only
222        """
223        # One-to-Many (classic)
224        one2many = []
225        for i in range(self.nl):
226            one2many.append(torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1))
227
228        # One-to-One (NMS-free)
229        one2one = []
230        for i in range(self.nl):
231            one2one.append(torch.cat((self.one2one_cv2[i](x[i]), self.one2one_cv3[i](x[i])), 1))
232
233        if self.training:
234            return one2many, one2one
235
236        # Inference path
237        if self.end2end:
238            # decode one2one → (B, 300, 6) style later
239            return self._inference(one2one)
240        else:
241            return self._inference(one2many)
242
243    def _inference(self, x):
244        """Simple decode for teaching (real code has more careful anchor decoding)."""
245        dbox = []
246        cls  = []
247        for i, feat in enumerate(x):
248            b, _, h, w = feat.shape
249            feat = feat.view(b, self.no, h * w)
250            box, score = feat.split((4, self.nc), 1)
251            dbox.append(box)
252            cls.append(score.sigmoid())
253        return torch.cat(dbox, 2), torch.cat(cls, 2)
254
255
256# =====================================================================
257# 6. Full YOLO26 Model
258# =====================================================================
259
260class YOLO26(nn.Module):
261    """
262    Educational YOLO26 architecture.
263    Scale: nano-ish (you can change width/depth multipliers).
264    """
265    def __init__(self, nc=80, width=0.25, depth=0.33):
266        super().__init__()
267        self.nc = nc
268
269        # channel helper
270        def c(x): return max(round(x * width), 8)
271        def n(x): return max(round(x * depth), 1)
272
273        # ---------------- Backbone ----------------
274        self.stem = Conv(3, c(64), 3, 2)                     # P1/2
275
276        self.stage1 = nn.Sequential(                         # P2/4
277            Conv(c(64), c(128), 3, 2),
278            C3k2(c(128), c(128), n=n(2), c3k=False),
279        )
280        self.stage2 = nn.Sequential(                         # P3/8
281            Conv(c(128), c(256), 3, 2),
282            C3k2(c(256), c(256), n=n(2), c3k=False),
283        )
284        self.stage3 = nn.Sequential(                         # P4/16
285            Conv(c(256), c(512), 3, 2),
286            C3k2(c(512), c(512), n=n(2), c3k=True),
287        )
288        self.stage4 = nn.Sequential(                         # P5/32
289            Conv(c(512), c(1024), 3, 2),
290            C3k2(c(1024), c(1024), n=n(2), c3k=True),
291            SPPF(c(1024), c(1024), k=5),
292            C2PSA(c(1024), c(1024), n=n(2)),                 # attention
293        )
294
295        # ---------------- Neck (FPN + PAN style) ----------------
296        self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
297
298        # Top-down
299        self.lateral4 = Conv(c(1024), c(512), 1, 1)
300        self.fusion4  = C3k2(c(1024), c(512), n=n(2), c3k=False)   # after concat P5↑ + P4
301
302        self.lateral3 = Conv(c(512), c(256), 1, 1)
303        self.fusion3  = C3k2(c(512), c(256), n=n(2), c3k=False)    # after concat P4↑ + P3
304
305        # Bottom-up
306        self.down3 = Conv(c(256), c(256), 3, 2)
307        self.fusion3b = C3k2(c(512), c(512), n=n(2), c3k=False)
308
309        self.down4 = Conv(c(512), c(512), 3, 2)
310        self.fusion4b = C3k2(c(1024), c(1024), n=n(2), c3k=True)
311
312        # ---------------- Head ----------------
313        self.detect = Detect(nc=nc, ch=(c(256), c(512), c(1024)))
314
315        self._init_weights()
316
317    def _init_weights(self):
318        for m in self.modules():
319            if isinstance(m, nn.Conv2d):
320                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
321            elif isinstance(m, nn.BatchNorm2d):
322                m.eps = 1e-3
323                m.momentum = 0.03
324
325    def forward(self, x):
326        # Backbone
327        x = self.stem(x)
328        x = self.stage1(x)
329        p3 = self.stage2(x)          # P3
330        p4 = self.stage3(p3)         # P4
331        p5 = self.stage4(p4)         # P5
332
333        # Neck – top-down
334        x = self.lateral4(p5)
335        x = self.upsample(x)
336        x = torch.cat([x, p4], 1)
337        p4_out = self.fusion4(x)
338
339        x = self.lateral3(p4_out)
340        x = self.upsample(x)
341        x = torch.cat([x, p3], 1)
342        p3_out = self.fusion3(x)
343
344        # Neck – bottom-up
345        x = self.down3(p3_out)
346        x = torch.cat([x, p4_out], 1)
347        p4_out = self.fusion3b(x)
348
349        x = self.down4(p4_out)
350        x = torch.cat([x, p5], 1)
351        p5_out = self.fusion4b(x)
352
353        # Head
354        return self.detect([p3_out, p4_out, p5_out])
355
356
357# =====================================================================
358# 7. Quick test / learning demo
359# =====================================================================
360
361if __name__ == "__main__":
362    device = "cuda" if torch.cuda.is_available() else "cpu"
363    model = YOLO26(nc=80, width=0.25, depth=0.33).to(device)   # nano-ish
364    model.eval()
365
366    x = torch.randn(1, 3, 640, 640, device=device)
367    with torch.no_grad():
368        out = model(x)
369
370    print("Input shape :", x.shape)
371    if isinstance(out, tuple):
372        print("Training mode outputs (one2many, one2one)")
373    else:
374        boxes, scores = out
375        print("Inference boxes shape :", boxes.shape)
376        print("Inference scores shape:", scores.shape)
377
378    # parameter count
379    n_params = sum(p.numel() for p in model.parameters())
380    print(f"Parameters: {n_params / 1e6:.2f} M")
381
382    # Show architecture summary
383    print("\n=== YOLO26 Educational Architecture ===")
384    print("Backbone : Stem → C3k2 stages → SPPF(residual) → C2PSA")
385    print("Neck     : FPN (top-down) + PAN (bottom-up) with C3k2")
386    print("Head     : Dual (One-to-Many + One-to-One), DFL-free")
387    print("Attention: C2PSA (Position-Sensitive Attention)")

How to read / learn this code

Recommended learning path

1. Run the __main__ block and print shapes at every stage.
2. Change width / depth and watch parameter count change.
3. Replace C2PSA with a pure Transformer encoder block and compare speed.
4. Implement a simple decode function that turns the one-to-one head into (B, 300, 6).

This gives you a complete, understandable YOLO26-style architecture focused on learning.