vision_rs/models/yolo/yolo26/blocks/
sppf.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::{
19    dtype::Float,
20    graph::{Op, SymTensor},
21    name_scope::name_scope,
22    nn::{Layer, pool::MaxPool2d},
23};
24
25use super::{concat::concat, conv::{conv, conv_bn}};
26
27fn elem_add(a: SymTensor, b: SymTensor) -> SymTensor {
28    let shape = a.shape.clone();
29    let node_id = a.graph.borrow_mut().add_node(
30        Op::Add, vec![a.node_id, b.node_id], a.dtype, shape.clone(),
31    );
32    SymTensor { node_id, graph: a.graph.clone(), dtype: a.dtype, shape }
33}
34
35/// Spatial Pyramid Pooling - Fast (SPPF).
36///
37/// Matches `ultralytics.nn.modules.block.SPPF(c1, c2, k=5, n=3, shortcut)`:
38///   cv1:      Conv(c1, c//2, 1, 1, act=False) — Conv+BN only, no SiLU
39///   pool:     MaxPool2d(k=5, stride=1, padding=2) applied 3 times
40///   cv2:      Conv(4 * c//2, c2, 1, 1) — Conv+BN+SiLU
41///   shortcut: residual add (output + input) when shortcut=true and c_in==c_out
42///             (YOLO26 passes shortcut=true; default is false)
43pub fn sppf<D: Float>(c_in: usize, c_out: usize, shortcut: bool) -> impl Fn(SymTensor) -> SymTensor {
44    let c = c_in / 2;
45    let add = shortcut && c_in == c_out;
46    let cv1 = conv_bn::<D>(c_in, c, 1, 1, 1);
47    let cv2 = conv::<D>(4 * c, c_out, 1, 1);
48    let pool = || {
49        MaxPool2d::<D, SymTensor, SymTensor, 4>::with_padding((5, 5), (1, 1), (2, 2))
50    };
51
52    move |x: SymTensor| {
53        let y   = { let _g = name_scope("cv1"); cv1(x.clone()) };
54        let p1  = pool().call(y.clone());
55        let p2  = pool().call(p1.clone());
56        let p3  = pool().call(p2.clone());
57        let cat = concat()(vec![y, p1, p2, p3]);
58        let out = { let _g = name_scope("cv2"); cv2(cat) };
59        if add { elem_add(out, x) } else { out }
60    }
61}
vision_rs/models/yolo/yolo26/blocks/sppf.rs

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