vision_rs/models/yolo/yolo26/blocks/
bottleneck.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::conv::conv;
21
22// ── Graph helper ──────────────────────────────────────────────────────────────
23
24pub(super) fn elem_add(a: SymTensor, b: SymTensor) -> SymTensor {
25    let shape = a.shape.clone();
26    let node_id = a.graph.borrow_mut().add_node(
27        Op::Add, vec![a.node_id, b.node_id], a.dtype, shape.clone(),
28    );
29    SymTensor { node_id, graph: a.graph.clone(), dtype: a.dtype, shape }
30}
31
32// ── Bottleneck variants ───────────────────────────────────────────────────────
33
34/// `conv(c → c//2, k=3) → conv(c//2 → c, k=3)`.
35///
36/// Matches `ultralytics.nn.modules.block.Bottleneck(k=(3,3), e=0.5)` defaults.
37pub fn bottleneck_std<D: Float + 'static>(c: usize, shortcut: bool) -> impl Fn(SymTensor) -> SymTensor {
38    let c_inner = (c as f32 * 0.5) as usize;
39    let cv1 = conv::<D>(c, c_inner, 3, 1);
40    let cv2 = conv::<D>(c_inner, c, 3, 1);
41    move |x: SymTensor| {
42        let y = {
43            let tmp = { let _g = name_scope("cv1"); cv1(x.clone()) };
44            let _g = name_scope("cv2"); cv2(tmp)
45        };
46        if shortcut { elem_add(x, y) } else { y }
47    }
48}
49
50/// `conv(c → c, k=3) → conv(c → c, k=3)`.
51///
52/// Used as the inner bottleneck inside `C3k` blocks (e=1.0).
53pub fn bottleneck_3x3<D: Float + 'static>(c: usize, shortcut: bool) -> impl Fn(SymTensor) -> SymTensor {
54    let cv1 = conv::<D>(c, c, 3, 1);
55    let cv2 = conv::<D>(c, c, 3, 1);
56    move |x: SymTensor| {
57        let y = {
58            let tmp = { let _g = name_scope("cv1"); cv1(x.clone()) };
59            let _g = name_scope("cv2"); cv2(tmp)
60        };
61        if shortcut { elem_add(x, y) } else { y }
62    }
63}
vision_rs/models/yolo/yolo26/blocks/bottleneck.rs

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