vision_rs/models/yolo/yolo26/blocks/
c3k2.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
18use teeny_core::{dtype::Float, graph::{Op, SymTensor}, name_scope::name_scope};
19
20use super::bottleneck::{bottleneck_3x3, bottleneck_std};
21use super::c2psa::psa_block;
22use super::conv::conv;
23
24// ── Graph helpers ─────────────────────────────────────────────────────────────
25
26fn channel_chunk(x: SymTensor, c_total: usize, chunk_c: usize, chunk_offset: usize) -> SymTensor {
27    let op = Op::ChannelChunk { c_total, chunk_c, chunk_offset };
28    let shape = vec![x.shape[0], Some(chunk_c), x.shape[2], x.shape[3]];
29    let node_id = x.graph.borrow_mut().add_node(op, vec![x.node_id], x.dtype, shape.clone());
30    SymTensor { node_id, graph: x.graph.clone(), dtype: x.dtype, shape }
31}
32
33fn channel_cat(tensors: Vec<SymTensor>, c_total: usize) -> SymTensor {
34    let first = &tensors[0];
35    let shape = vec![first.shape[0], Some(c_total), first.shape[2], first.shape[3]];
36    let inputs: Vec<usize> = tensors.iter().map(|t| t.node_id).collect();
37    let node_id = first.graph.borrow_mut().add_node(
38        Op::ChannelCat { c_total }, inputs, first.dtype, shape.clone(),
39    );
40    SymTensor { node_id, graph: first.graph.clone(), dtype: first.dtype, shape }
41}
42
43// ── C3k inner block ───────────────────────────────────────────────────────────
44
45/// Matches ultralytics C3k(c, c, n=2, shortcut, e=0.5).
46fn c3k_inner<D: Float + 'static>(c: usize, shortcut: bool) -> impl Fn(SymTensor) -> SymTensor {
47    let c_h = ((c as f32) * 0.5) as usize;
48    let cv1 = conv::<D>(c, c_h, 1, 1);
49    let cv2 = conv::<D>(c, c_h, 1, 1);
50    let cv3 = conv::<D>(2 * c_h, c, 1, 1);
51    let m0 = bottleneck_3x3::<D>(c_h, shortcut);
52    let m1 = bottleneck_3x3::<D>(c_h, shortcut);
53    move |x: SymTensor| {
54        let after_cv1 = { let _g = name_scope("cv1"); cv1(x.clone()) };
55        let after_m0  = { let _g = name_scope("m.0"); m0(after_cv1) };
56        let after_m1  = { let _g = name_scope("m.1"); m1(after_m0) };
57        let after_cv2 = { let _g = name_scope("cv2"); cv2(x) };
58        let _g = name_scope("cv3");
59        cv3(channel_cat(vec![after_m1, after_cv2], 2 * c_h))
60    }
61}
62
63/// Sequential([Bottleneck, PSABlock]) inner for model.22.
64///
65/// Matches the special C3k2 variant used at P5/32 in YOLO26.
66/// Named as "0" (Bottleneck) and "1" (PSABlock) — ultralytics Sequential indexing.
67fn bottleneck_psa_seq<D: Float + Send + Sync + 'static>(
68    c: usize,
69    shortcut: bool,
70    num_heads: usize,
71    key_dim: usize,
72) -> impl Fn(SymTensor) -> SymTensor {
73    let bneck = bottleneck_std::<D>(c, shortcut);
74    let psa   = psa_block::<D>(c, num_heads, key_dim);
75    move |x: SymTensor| {
76        let x = { let _g = name_scope("0"); bneck(x) };
77        { let _g = name_scope("1"); psa(x) }
78    }
79}
80
81// ── C3k2 block ────────────────────────────────────────────────────────────────
82
83/// C3k2: the primary feature-extraction block in YOLO11/26.
84///
85/// Matches `ultralytics.nn.modules.block.C3k2` (which inherits from C2f).
86/// Forward:
87///   h = cv1(x)                          // [N, 2*c, H, W]
88///   [y0, y1] = h.chunk(2, dim=1)        // each [N, c, H, W]
89///   parts = [y0, y1]
90///   last = y1
91///   for each bottleneck b:
92///       last = b(last); parts.append(last)
93///   return cv2(cat(parts))              // [(2+n)*c → c_out]
94pub fn c3k2<D: Float + 'static>(
95    c_in: usize,
96    c_out: usize,
97    n: usize,
98    c3k: bool,
99    shortcut: bool,
100    e: f32,
101) -> impl Fn(SymTensor) -> SymTensor {
102    let c = (c_out as f32 * e) as usize;
103    let cv1 = conv::<D>(c_in, 2 * c, 1, 1);
104    let cv2 = conv::<D>((2 + n) * c, c_out, 1, 1);
105    let bottlenecks: Vec<Box<dyn Fn(SymTensor) -> SymTensor>> = (0..n)
106        .map(|_| -> Box<dyn Fn(SymTensor) -> SymTensor> {
107            if c3k {
108                Box::new(c3k_inner::<D>(c, shortcut))
109            } else {
110                Box::new(bottleneck_std::<D>(c, shortcut))
111            }
112        })
113        .collect();
114    move |x: SymTensor| {
115        let h = { let _g = name_scope("cv1"); cv1(x) };
116        let y0 = channel_chunk(h.clone(), 2 * c, c, 0);
117        let y1 = channel_chunk(h, 2 * c, c, c);
118        let mut last = y1.clone();
119        let mut parts = vec![y0, y1];
120        for (i, b) in bottlenecks.iter().enumerate() {
121            let _g = name_scope(format!("m.{i}"));
122            last = b(last);
123            parts.push(last.clone());
124        }
125        { let _g = name_scope("cv2"); cv2(channel_cat(parts, (2 + n) * c)) }
126    }
127}
128
129/// C3k2 variant whose inner blocks are `Sequential([Bottleneck, PSABlock])`.
130///
131/// Used for model.22 (P5/32 detect-path block) in YOLO26. The inner channel
132/// width `c = c_out * e`; `num_heads = c / 64`, `key_dim = 32` follow the
133/// same convention as C2PSA.
134pub fn c3k2_psa<D: Float + Send + Sync + 'static>(
135    c_in: usize,
136    c_out: usize,
137    n: usize,
138    shortcut: bool,
139    e: f32,
140) -> impl Fn(SymTensor) -> SymTensor {
141    let c = (c_out as f32 * e) as usize;
142    let num_heads = c / 64;
143    let key_dim = 32;
144    let cv1 = conv::<D>(c_in, 2 * c, 1, 1);
145    let cv2 = conv::<D>((2 + n) * c, c_out, 1, 1);
146    let bottlenecks: Vec<Box<dyn Fn(SymTensor) -> SymTensor>> = (0..n)
147        .map(|_| -> Box<dyn Fn(SymTensor) -> SymTensor> {
148            Box::new(bottleneck_psa_seq::<D>(c, shortcut, num_heads, key_dim))
149        })
150        .collect();
151    move |x: SymTensor| {
152        let h = { let _g = name_scope("cv1"); cv1(x) };
153        let y0 = channel_chunk(h.clone(), 2 * c, c, 0);
154        let y1 = channel_chunk(h, 2 * c, c, c);
155        let mut last = y1.clone();
156        let mut parts = vec![y0, y1];
157        for (i, b) in bottlenecks.iter().enumerate() {
158            let _g = name_scope(format!("m.{i}"));
159            last = b(last);
160            parts.push(last.clone());
161        }
162        { let _g = name_scope("cv2"); cv2(channel_cat(parts, (2 + n) * c)) }
163    }
164}
vision_rs/models/yolo/yolo26/blocks/c3k2.rs

vision_rs/models/yolo/yolo26/blocks/
c3k2.rs
