vision_rs/models/yolo/yolo26/blocks/
conv.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::SymTensor,
21    name_scope::name_scope,
22    nn::{Layer, activation::sigmoid::Silu, batchnorm::BatchNorm2d, conv2d::Conv2d},
23};
24
25/// Conv2d → BatchNorm2d → SiLU — the basic building block of YOLO26.
26///
27/// Matches `ultralytics.nn.modules.conv.Conv(c1, c2, k, s)`:
28///   - bias=False (BN subsumes the bias term)
29///   - autopad: p = k / 2 (same-padding for odd kernels)
30///   - BN eps=0.001 matching the pretrained YOLO26n weights (trained with eps=1e-3)
31pub fn conv<D: Float>(
32    c_in: usize,
33    c_out: usize,
34    k: usize,
35    s: usize,
36) -> impl Fn(SymTensor) -> SymTensor {
37    let p = k / 2;
38    let conv2d = Conv2d::<D, SymTensor, SymTensor, 4>::new(c_in, c_out, (k, k), (s, s), (p, p), false);
39    let bn = BatchNorm2d::<D, SymTensor, SymTensor, 4>::new(c_out).with_eps(0.001);
40    let act = Silu::<D, SymTensor, 4>::new();
41    move |x: SymTensor| {
42        let x = { let _g = name_scope("conv"); conv2d.call(x) };
43        let x = { let _g = name_scope("bn"); bn.call(x) };
44        act.call(x)
45    }
46}
47
48/// Conv2d(groups) → BatchNorm2d — no activation.
49///
50/// Used for layers in PSABlock that skip the SiLU activation (qkv, pe, proj, ffn1).
51/// `g=1` for standard 1×1 convs; `g=c_in=c_out` for depthwise.
52/// BN eps=0.001 matching the pretrained YOLO26n weights.
53pub fn conv_bn<D: Float>(
54    c_in: usize,
55    c_out: usize,
56    k: usize,
57    s: usize,
58    g: usize,
59) -> impl Fn(SymTensor) -> SymTensor {
60    let p = k / 2;
61    let conv2d = Conv2d::<D, SymTensor, SymTensor, 4>::new_grouped(
62        c_in, c_out, (k, k), (s, s), (p, p), false, g,
63    );
64    let bn = BatchNorm2d::<D, SymTensor, SymTensor, 4>::new(c_out).with_eps(0.001);
65    move |x: SymTensor| {
66        let x = { let _g = name_scope("conv"); conv2d.call(x) };
67        { let _g = name_scope("bn"); bn.call(x) }
68    }
69}
70
71/// Depth-wise Conv2d (groups = c) → BatchNorm2d → SiLU.
72///
73/// Matches `ultralytics DWConv(c, c, k)` with act=True.
74/// BN eps=0.001 matching the pretrained YOLO26n weights.
75pub fn dwconv<D: Float>(c: usize, k: usize, s: usize) -> impl Fn(SymTensor) -> SymTensor {
76    let p = k / 2;
77    let conv2d = Conv2d::<D, SymTensor, SymTensor, 4>::new_grouped(c, c, (k, k), (s, s), (p, p), false, c);
78    let bn = BatchNorm2d::<D, SymTensor, SymTensor, 4>::new(c).with_eps(0.001);
79    let act = Silu::<D, SymTensor, 4>::new();
80    move |x: SymTensor| {
81        let x = { let _g = name_scope("conv"); conv2d.call(x) };
82        let x = { let _g = name_scope("bn"); bn.call(x) };
83        act.call(x)
84    }
85}
86
87/// Plain Conv2d with bias — no BatchNorm, no activation.
88///
89/// Matches `nn.Conv2d(c_in, c_out, k, bias=True)` from the ultralytics Detect head.
90pub fn conv_plain<D: Float>(
91    c_in: usize,
92    c_out: usize,
93    k: usize,
94    s: usize,
95) -> impl Fn(SymTensor) -> SymTensor {
96    let p = k / 2;
97    let layer = Conv2d::<D, SymTensor, SymTensor, 4>::new(c_in, c_out, (k, k), (s, s), (p, p), true);
98    move |x| layer.call(x)
99}
vision_rs/models/yolo/yolo26/blocks/conv.rs

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