vision_rs/models/yolo/yolo26/
mod.rs

1/*
2 * Copyright 2026 Teenygrad
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17
18//! YOLO26 detection model.
19//!
20//! Architecture reference: `ultralytics/cfg/models/26/yolo26.yaml`.
21//!
22//! Returns raw `DetectOutput { boxes (B, 4·A), scores (B, nc·A) }` — the
23//! training-mode layout.  For inference, apply `detect_decode` on `boxes`
24//! using the anchor grid and strides [8, 16, 32] appropriate for the
25//! runtime input resolution.
26
27use teeny_core::{dtype::Float, graph::SymTensor, name_scope::name_scope};
28
29use blocks::{
30    c2psa::c2psa,
31    c3k2::{c3k2, c3k2_psa},
32    concat::concat,
33    conv::conv,
34    detect::{DetectHead, DetectOutput, DualDetectOutput, detect},
35    sppf::sppf,
36    upsample::upsample,
37};
38
39/// Building blocks (backbone/neck/head layers) used to assemble YOLO26.
40pub mod blocks;
41
42// ── Variant + Config ──────────────────────────────────────────────────────────
43
44/// Per-variant scaling parameters derived from the yaml model definition.
45#[derive(Clone, Copy, Debug)]
46pub struct Yolo26Config {
47    /// Depth multiplier, scaling block repeat counts.
48    pub depth: f32,
49    /// Width multiplier, scaling channel counts.
50    pub width: f32,
51    /// Maximum channel count cap.
52    pub mc: usize,
53}
54
55/// YOLO26 model size variant.
56#[derive(Clone, Copy, Debug)]
57pub enum Yolo26Variant {
58    /// Nano.
59    N,
60    /// Small.
61    S,
62    /// Medium.
63    M,
64    /// Large.
65    L,
66    /// Extra-large.
67    XL,
68}
69
70impl Yolo26Variant {
71    /// Returns this variant's depth/width/max-channel scaling parameters.
72    pub fn config(&self) -> Yolo26Config {
73        match self {
74            Yolo26Variant::N => Yolo26Config {
75                depth: 0.5,
76                width: 0.25,
77                mc: 1024,
78            },
79            Yolo26Variant::S => Yolo26Config {
80                depth: 0.5,
81                width: 0.50,
82                mc: 1024,
83            },
84            Yolo26Variant::M => Yolo26Config {
85                depth: 0.5,
86                width: 1.00,
87                mc: 512,
88            },
89            Yolo26Variant::L => Yolo26Config {
90                depth: 1.0,
91                width: 1.00,
92                mc: 512,
93            },
94            Yolo26Variant::XL => Yolo26Config {
95                depth: 1.0,
96                width: 1.50,
97                mc: 512,
98            },
99        }
100    }
101}
102
103// ── Scaling helpers ───────────────────────────────────────────────────────────
104
105/// Scale a yaml base channel count by width, capped at max_channels.
106fn ch(base: usize, width: f32, mc: usize) -> usize {
107    (base.min(mc) as f32 * width).round() as usize
108}
109
110/// Scale a yaml repeat count by depth, minimum 1.
111fn rep(base: usize, depth: f32) -> usize {
112    if base > 1 {
113        ((base as f32 * depth).max(1.0)).round() as usize
114    } else {
115        base
116    }
117}
118
119// ── Shared backbone + FPN neck ────────────────────────────────────────────────
120
121/// Build the YOLO26 backbone and FPN neck layers.
122///
123/// Returns `(neck_fn, [c2, c3, c4])` where `neck_fn` maps an input image
124/// tensor to the three FPN feature maps `(p3d, p4d, p5d)` consumed by the
125/// detect head(s), and `[c2, c3, c4]` are the corresponding channel widths.
126fn build_neck<D: Float + Send + Sync + 'static>(
127    variant: &Yolo26Variant,
128) -> (impl Fn(SymTensor) -> (SymTensor, SymTensor, SymTensor), usize, usize, usize) {
129    let cfg = variant.config();
130    let (d, w, mc) = (cfg.depth, cfg.width, cfg.mc);
131
132    let c0 = ch(64, w, mc);
133    let c1 = ch(128, w, mc);
134    let c2 = ch(256, w, mc);
135    let c3 = ch(512, w, mc);
136    let c4 = ch(1024, w, mc);
137    let n = rep(2, d);
138
139    // Backbone layers (yaml 0–10)
140    let l0  = conv::<D>(3, c0, 3, 2);
141    let l1  = conv::<D>(c0, c1, 3, 2);
142    let l2  = c3k2::<D>(c1, c2, n, false, true, 0.25);
143    let l3  = conv::<D>(c2, c2, 3, 2);
144    let l4  = c3k2::<D>(c2, c3, n, false, true, 0.25);
145    let l5  = conv::<D>(c3, c3, 3, 2);
146    let l6  = c3k2::<D>(c3, c3, n, true, true, 0.5);
147    let l7  = conv::<D>(c3, c4, 3, 2);
148    let l8  = c3k2::<D>(c4, c4, n, true, true, 0.5);
149    let l9  = sppf::<D>(c4, c4, true);
150    let l10 = c2psa::<D>(c4, c4, n, 0.5);
151
152    // FPN neck layers (yaml 11–22)
153    let up  = upsample(2, 2);
154    let cat = concat();
155    let l13 = c3k2::<D>(c4 + c3, c3, n, true, true, 0.5);
156    let l16 = c3k2::<D>(c3 + c3, c2, n, true, true, 0.5);
157    let l17 = conv::<D>(c2, c2, 3, 2);
158    let l19 = c3k2::<D>(c2 + c3, c3, n, true, true, 0.5);
159    let l20 = conv::<D>(c3, c3, 3, 2);
160    let l22 = c3k2_psa::<D>(c3 + c4, c4, 1, true, 0.5);
161
162    let neck_fn = move |x: SymTensor| -> (SymTensor, SymTensor, SymTensor) {
163        let x = { let _g = name_scope("model.0");  l0(x) };
164        let x = { let _g = name_scope("model.1");  l1(x) };
165        let x = { let _g = name_scope("model.2");  l2(x) };
166        let x = { let _g = name_scope("model.3");  l3(x) };
167        let p3 = { let _g = name_scope("model.4");  l4(x) };
168        let x  = { let _g = name_scope("model.5");  l5(p3.clone()) };
169        let p4 = { let _g = name_scope("model.6");  l6(x) };
170        let x  = { let _g = name_scope("model.7");  l7(p4.clone()) };
171        let x  = { let _g = name_scope("model.8");  l8(x) };
172        let x  = { let _g = name_scope("model.9");  l9(x) };
173        let p5 = { let _g = name_scope("model.10"); l10(x) };
174
175        let x   = up(p5.clone());
176        let x   = cat(vec![x, p4]);
177        let nk4 = { let _g = name_scope("model.13"); l13(x) };
178
179        let x   = up(nk4.clone());
180        let x   = cat(vec![x, p3]);
181        let p3d = { let _g = name_scope("model.16"); l16(x) };
182
183        let x   = { let _g = name_scope("model.17"); l17(p3d.clone()) };
184        let x   = cat(vec![x, nk4]);
185        let p4d = { let _g = name_scope("model.19"); l19(x) };
186
187        let x   = { let _g = name_scope("model.20"); l20(p4d.clone()) };
188        let x   = cat(vec![x, p5]);
189        let p5d = { let _g = name_scope("model.22"); l22(x) };
190
191        (p3d, p4d, p5d)
192    };
193
194    (neck_fn, c2, c3, c4)
195}
196
197// ── Public model constructors ─────────────────────────────────────────────────
198
199/// YOLO26 single-head forward closure (inference or one-head training).
200///
201/// # Arguments
202/// * `nc`      — number of detection classes (e.g. 80 for COCO)
203/// * `variant` — model size variant (N / S / M / L / XL)
204/// * `head`    — `OneToMany` (cv2/cv3, dense) or `OneToOne` (one2one_cv2/cv3, inference)
205///
206/// For dual-assignment training use [`yolo26_dual`] instead.
207pub fn yolo26<D: Float + Send + Sync + 'static>(
208    nc: usize,
209    variant: &Yolo26Variant,
210    head: DetectHead,
211) -> impl Fn(SymTensor) -> DetectOutput {
212    let (neck, c2, c3, c4) = build_neck::<D>(variant);
213    let head_fn = detect::<D>(nc, &[c2, c3, c4], head);
214
215    move |x: SymTensor| {
216        let (p3d, p4d, p5d) = neck(x);
217        { let _g = name_scope("model.23"); head_fn(vec![p3d, p4d, p5d]) }
218    }
219}
220
221/// YOLO26 dual-head forward closure for training with consistent dual assignment.
222///
223/// Traces **both** the one2many head (cv2/cv3, TAL top_k=10) and the one2one
224/// head (one2one_cv2/cv3, TAL top_k=1) in a single graph, sharing the backbone
225/// and FPN neck.  The returned [`DualDetectOutput`] exposes both sets of
226/// predictions so the caller can compute weighted losses for each.
227///
228/// Loss weighting schedule (ultralytics-style): early in training weight the
229/// one2many head more heavily (it provides dense, stable gradients); gradually
230/// shift weight toward the one2one head so it matches inference behaviour by the
231/// end of training.  Use `Yolo26Loss::compute_grads_dual` (in
232/// `models::yolo::loss::yolo26`, behind the `cuda` feature) to apply this.
233pub fn yolo26_dual<D: Float + Send + Sync + 'static>(
234    nc: usize,
235    variant: &Yolo26Variant,
236) -> impl Fn(SymTensor) -> DualDetectOutput {
237    let (neck, c2, c3, c4) = build_neck::<D>(variant);
238    let head_o2m = detect::<D>(nc, &[c2, c3, c4], DetectHead::OneToMany);
239    let head_o2o = detect::<D>(nc, &[c2, c3, c4], DetectHead::OneToOne);
240
241    move |x: SymTensor| {
242        let (p3d, p4d, p5d) = neck(x);
243        // one2many is traced before one2one so its terminal nodes get lower DAG
244        // indices — the training loop relies on this for stable node identification.
245        let one2many = {
246            let _g = name_scope("model.23");
247            head_o2m(vec![p3d.clone(), p4d.clone(), p5d.clone()])
248        };
249        let one2one = {
250            let _g = name_scope("model.23");
251            head_o2o(vec![p3d, p4d, p5d])
252        };
253        DualDetectOutput { one2many, one2one }
254    }
255}
vision_rs/models/yolo/yolo26/mod.rs

vision_rs/models/yolo/yolo26/
mod.rs
