vision_rs/models/yolo/loss/
yolo26.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 training loss: CIoU box regression + BCE classification.
19//!
20//! # Pipeline (per training step)
21//!
22//! 1. **Decode** raw LTRB predictions to XYWH using the anchor grid.
23//! 2. **Assign** GT boxes to anchors (CPU-side TaskAlignedAssigner).
24//! 3. **Box loss**: CIoU forward (saved activations) then backward.
25//! 4. **Cls loss**: BCE cls backward on all anchors.
26//! 5. **Decode backward**: d_xywh → d_ltrb through the linear decode.
27//! 6. Return `d_boxes` `[B, 4*A]` and `d_scores` `[B, nc*A]` (host f32).
28
29pub use super::{anchor::AnchorGrid, assign::TaskAlignedAssigner};
30
31/// CPU reference CIoU forward — returns saved activations (iou, v, alpha).
32pub fn ciou_forward_cpu(pred: &[f32], target: &[f32], n: usize)
33    -> (Vec<f32>, Vec<f32>, Vec<f32>)
34{
35    const EPS: f32 = 1e-7;
36    let pi2 = std::f32::consts::PI * std::f32::consts::PI;
37    let mut iou   = vec![0.0f32; n];
38    let mut v_vec = vec![0.0f32; n];
39    let mut alpha = vec![0.0f32; n];
40
41    for i in 0..n {
42        let px = pred[i];           let py = pred[n + i];
43        let pw = pred[2 * n + i];   let ph = pred[3 * n + i];
44        let tx = target[i];         let ty = target[n + i];
45        let tw = target[2 * n + i]; let th = target[3 * n + i];
46
47        let px1 = px - pw * 0.5; let px2 = px + pw * 0.5;
48        let py1 = py - ph * 0.5; let py2 = py + ph * 0.5;
49        let tx1 = tx - tw * 0.5; let tx2 = tx + tw * 0.5;
50        let ty1 = ty - th * 0.5; let ty2 = ty + th * 0.5;
51
52        let inter = ((px2.min(tx2) - px1.max(tx1)).max(0.0))
53                  * ((py2.min(ty2) - py1.max(ty1)).max(0.0));
54        let union = pw * ph + tw * th - inter;
55        let iou_i = inter / (union + EPS);
56
57        let v_i = (4.0 / pi2)
58            * ((tw / (th + EPS)).atan() - (pw / (ph + EPS)).atan()).powi(2);
59        let alpha_i = v_i / (1.0 - iou_i + v_i + EPS);
60
61        iou[i] = iou_i; v_vec[i] = v_i; alpha[i] = alpha_i;
62    }
63    (iou, v_vec, alpha)
64}
65
66// ── CUDA-only training loss ───────────────────────────────────────────────────
67
68#[cfg(feature = "cuda")]
69mod cuda_impl {
70    use super::*;
71    use teeny_compiler::compiler::{driver::cuda::compile_kernel, target::cuda::Target};
72    use teeny_core::device::{Device, buffer::Buffer};
73    use teeny_cuda::{compiler::target::Capability, device::CudaDevice, errors::Result, testing};
74    use crate::models::yolo::kernels::loss::{
75        ciou::{YoloCiouLossBackward, YoloCiouLossForward},
76        cls::YoloBceClsLossBackward,
77    };
78
79    /// CUDA training loss for YOLO26: computes d_boxes and d_scores.
80    pub struct Yolo26Loss {
81        /// Anchor grid (points + strides) for the model's input resolution.
82        pub grid:         AnchorGrid,
83        /// One2many assigner (top_k=10) — used during `compute_grads` and the o2m head of `compute_grads_dual`.
84        pub assigner:     TaskAlignedAssigner,
85        /// One2one assigner (top_k=1) — used for the o2o head of `compute_grads_dual`.
86        pub assigner_o2o: TaskAlignedAssigner,
87        /// Number of detection classes.
88        pub nc:           usize,
89        /// Target GPU compute capability for compiled kernels.
90        pub cap:          Capability,
91        block_n:          i32,
92    }
93
94    impl Yolo26Loss {
95        /// Builds the loss state for a model trained at `img_h × img_w` with `nc` classes.
96        pub fn new(img_h: usize, img_w: usize, nc: usize, cap: Capability) -> Self {
97            Self {
98                grid:         AnchorGrid::yolo26(img_h, img_w),
99                assigner:     TaskAlignedAssigner::default(),
100                assigner_o2o: TaskAlignedAssigner { top_k: 1, ..TaskAlignedAssigner::default() },
101                nc, cap, block_n: 64,
102            }
103        }
104
105        /// Compute loss gradients for one batch (one2many head only).
106        ///
107        /// `boxes`   – raw LTRB predictions `[B, 4*A]` (host f32, channels-first)
108        /// `scores`  – class logits `[B, nc*A]` (host f32, channels-first)
109        ///
110        /// Returns `(d_boxes [B,4*A], d_scores [B,nc*A])` as host f32.
111        pub fn compute_grads(
112            &self,
113            device: &CudaDevice<'_>,
114            boxes:      &[f32],
115            scores:     &[f32],
116            gt_boxes_b: &[Vec<[f32; 4]>],
117            gt_cls_b:   &[Vec<usize>],
118        ) -> anyhow::Result<(Vec<f32>, Vec<f32>)> {
119            let target = Target::new(self.cap);
120            let bn = self.block_n;
121            let ciou_fwd_ptx = std::fs::read(compile_kernel(&YoloCiouLossForward::<f32>::new(bn), &target, true)?)?;
122            let ciou_bwd_ptx = std::fs::read(compile_kernel(&YoloCiouLossBackward::<f32>::new(bn), &target, true)?)?;
123            let cls_bwd_ptx  = std::fs::read(compile_kernel(&YoloBceClsLossBackward::<f32>::new(bn), &target, true)?)?;
124            let prog_ciou_fwd = testing::load_program_from_ptx::<YoloCiouLossForward<f32>>(&ciou_fwd_ptx)?;
125            let prog_ciou_bwd = testing::load_program_from_ptx::<YoloCiouLossBackward<f32>>(&ciou_bwd_ptx)?;
126            let prog_cls_bwd  = testing::load_program_from_ptx::<YoloBceClsLossBackward<f32>>(&cls_bwd_ptx)?;
127            self.compute_grads_for_head(
128                device, &prog_ciou_fwd, &prog_ciou_bwd, &prog_cls_bwd,
129                &self.assigner, boxes, scores, gt_boxes_b, gt_cls_b, 1.0,
130            )
131        }
132
133        /// Compute dual-head loss gradients for consistent assignment training.
134        ///
135        /// Runs TAL assignment independently for both heads using their respective
136        /// assigners (top_k=10 for o2m, top_k=1 for o2o), then scales the resulting
137        /// gradients by `w_o2m` and `w_o2o` respectively.
138        ///
139        /// # Loss weight schedule (ultralytics-style)
140        /// - `w_o2m = 1.0` (constant throughout training)
141        /// - `w_o2o = step / total_steps` (ramps 0→1 linearly; caller controls the schedule)
142        ///
143        /// Returns `(d_boxes_o2m, d_scores_o2m, d_boxes_o2o, d_scores_o2o)`.
144        #[allow(clippy::too_many_arguments, clippy::type_complexity)]
145        pub fn compute_grads_dual(
146            &self,
147            device:      &CudaDevice<'_>,
148            boxes_o2m:   &[f32],
149            scores_o2m:  &[f32],
150            boxes_o2o:   &[f32],
151            scores_o2o:  &[f32],
152            gt_boxes_b:  &[Vec<[f32; 4]>],
153            gt_cls_b:    &[Vec<usize>],
154            w_o2m:       f32,
155            w_o2o:       f32,
156        ) -> anyhow::Result<(Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>)> {
157            let target = Target::new(self.cap);
158            let bn = self.block_n;
159            let ciou_fwd_ptx = std::fs::read(compile_kernel(&YoloCiouLossForward::<f32>::new(bn), &target, true)?)?;
160            let ciou_bwd_ptx = std::fs::read(compile_kernel(&YoloCiouLossBackward::<f32>::new(bn), &target, true)?)?;
161            let cls_bwd_ptx  = std::fs::read(compile_kernel(&YoloBceClsLossBackward::<f32>::new(bn), &target, true)?)?;
162            let prog_ciou_fwd = testing::load_program_from_ptx::<YoloCiouLossForward<f32>>(&ciou_fwd_ptx)?;
163            let prog_ciou_bwd = testing::load_program_from_ptx::<YoloCiouLossBackward<f32>>(&ciou_bwd_ptx)?;
164            let prog_cls_bwd  = testing::load_program_from_ptx::<YoloBceClsLossBackward<f32>>(&cls_bwd_ptx)?;
165
166            let (d_boxes_o2m, d_scores_o2m) = self.compute_grads_for_head(
167                device, &prog_ciou_fwd, &prog_ciou_bwd, &prog_cls_bwd,
168                &self.assigner, boxes_o2m, scores_o2m, gt_boxes_b, gt_cls_b, w_o2m,
169            )?;
170            let (d_boxes_o2o, d_scores_o2o) = self.compute_grads_for_head(
171                device, &prog_ciou_fwd, &prog_ciou_bwd, &prog_cls_bwd,
172                &self.assigner_o2o, boxes_o2o, scores_o2o, gt_boxes_b, gt_cls_b, w_o2o,
173            )?;
174            Ok((d_boxes_o2m, d_scores_o2m, d_boxes_o2o, d_scores_o2o))
175        }
176
177        #[allow(clippy::too_many_arguments)]
178        fn compute_grads_for_head(
179            &self,
180            device:       &CudaDevice<'_>,
181            prog_ciou_fwd: &teeny_cuda::device::program::CudaProgram<'_, YoloCiouLossForward<f32>>,
182            prog_ciou_bwd: &teeny_cuda::device::program::CudaProgram<'_, YoloCiouLossBackward<f32>>,
183            prog_cls_bwd:  &teeny_cuda::device::program::CudaProgram<'_, YoloBceClsLossBackward<f32>>,
184            assigner:     &TaskAlignedAssigner,
185            boxes:        &[f32],
186            scores:       &[f32],
187            gt_boxes_b:   &[Vec<[f32; 4]>],
188            gt_cls_b:     &[Vec<usize>],
189            loss_weight:  f32,
190        ) -> anyhow::Result<(Vec<f32>, Vec<f32>)> {
191            let b  = gt_boxes_b.len();
192            let a  = self.grid.n_anchors;
193            let nc = self.nc;
194
195            assert_eq!(boxes.len(),  b * 4 * a);
196            assert_eq!(scores.len(), b * nc * a);
197
198            let mut d_boxes_out  = vec![0.0f32; b * 4 * a];
199            let mut d_scores_out = vec![0.0f32; b * nc * a];
200
201            for bi in 0..b {
202                let boxes_i  = &boxes[bi * 4 * a .. (bi + 1) * 4 * a];
203                let scores_i = &scores[bi * nc * a .. (bi + 1) * nc * a];
204
205                // 1. Decode LTRB → XYWH (CPU).
206                let xywh = self.grid.decode_ltrb_to_xywh(boxes_i);
207
208                // 2. Assign GT targets (CPU).
209                let assign = assigner.assign(
210                    &xywh, scores_i,
211                    &self.grid.cx, &self.grid.cy,
212                    &gt_boxes_b[bi], &gt_cls_b[bi],
213                );
214                let n_pos: usize = assign.is_positive.iter().filter(|&&p| p).count();
215                // Normalise by sum of soft targets, matching ultralytics target_scores_sum.
216                let target_scores_sum: f32 = (0..a)
217                    .filter(|&i| assign.is_positive[i])
218                    .map(|i| assign.soft_target[i])
219                    .sum::<f32>()
220                    .max(1.0);
221                let norm = loss_weight / target_scores_sum;
222
223                // 3. CIoU backward for positive anchors.
224                let mut d_xywh = vec![0.0f32; 4 * a];
225                if n_pos > 0 {
226                    let pos_idx: Vec<usize> = (0..a).filter(|&i| assign.is_positive[i]).collect();
227                    let np = pos_idx.len();
228
229                    let mut pred_pos   = vec![0.0f32; 4 * np];
230                    let mut target_pos = vec![0.0f32; 4 * np];
231                    for (j, &i) in pos_idx.iter().enumerate() {
232                        for ch in 0..4 {
233                            pred_pos[ch * np + j]   = xywh[ch * a + i];
234                            target_pos[ch * np + j] = assign.target_boxes[ch * a + i];
235                        }
236                    }
237
238                    // CIoU forward on GPU to get saved activations.
239                    let (iou, v, alpha) = self.ciou_fwd_gpu(
240                        device, prog_ciou_fwd, &pred_pos, &target_pos, np
241                    )?;
242
243                    // Box loss weight = soft_target (alignment-metric based), not raw IoU.
244                    let dy_pos: Vec<f32> = pos_idx.iter()
245                        .map(|&i| norm * assign.soft_target[i])
246                        .collect();
247
248                    let d_pred_pos = self.ciou_bwd_gpu(
249                        device, prog_ciou_bwd,
250                        &dy_pos, &pred_pos, &target_pos, &iou, &v, &alpha, np,
251                    )?;
252
253                    for (j, &i) in pos_idx.iter().enumerate() {
254                        for ch in 0..4 { d_xywh[ch * a + i] = d_pred_pos[ch * np + j]; }
255                    }
256                }
257
258                // 4. Decode backward: d_xywh → d_ltrb.
259                let d_ltrb = self.grid.decode_backward(&d_xywh);
260                d_boxes_out[bi * 4 * a .. (bi + 1) * 4 * a].copy_from_slice(&d_ltrb);
261
262                // 5. BCE cls backward — all anchors with soft targets.
263                //
264                // Soft labels: positive anchor's class slot gets soft_target (∈ (0,max_iou])
265                // rather than a hard 1.0.  All other slots (and all negatives) stay 0.
266                let mut cls_target = vec![0.0f32; nc * a];
267                for (i, &pos) in assign.is_positive.iter().enumerate() {
268                    if pos {
269                        let c = assign.target_cls[i];
270                        if c < nc { cls_target[c * a + i] = assign.soft_target[i]; }
271                    }
272                }
273                let dy_cls = vec![norm; a];
274
275                let d_scores_i = self.cls_bwd_gpu(
276                    device, prog_cls_bwd,
277                    &dy_cls, scores_i, &cls_target, a, nc,
278                )?;
279                d_scores_out[bi * nc * a .. (bi + 1) * nc * a].copy_from_slice(&d_scores_i);
280            }
281
282            Ok((d_boxes_out, d_scores_out))
283        }
284
285        // ── GPU helpers ───────────────────────────────────────────────────────
286
287        fn ciou_fwd_gpu(
288            &self, device: &CudaDevice<'_>,
289            prog: &teeny_cuda::device::program::CudaProgram<'_, YoloCiouLossForward<f32>>,
290            pred: &[f32], target: &[f32], n: usize,
291        ) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>)> {
292            let mut pred_buf   = device.buffer::<f32>(4 * n)?;
293            let mut target_buf = device.buffer::<f32>(4 * n)?;
294            let loss_buf   = device.buffer::<f32>(n)?;
295            let iou_buf    = device.buffer::<f32>(n)?;
296            let v_buf      = device.buffer::<f32>(n)?;
297            let alpha_buf  = device.buffer::<f32>(n)?;
298            pred_buf.to_device(pred)?;
299            target_buf.to_device(target)?;
300            let cfg = testing::launch_config_with_grid(n.div_ceil(self.block_n as usize), prog);
301            device.launch(prog, &cfg, (
302                pred_buf.as_device_ptr()   as *mut f32,
303                target_buf.as_device_ptr() as *mut f32,
304                loss_buf.as_device_ptr()   as *mut f32,
305                iou_buf.as_device_ptr()    as *mut f32,
306                v_buf.as_device_ptr()      as *mut f32,
307                alpha_buf.as_device_ptr()  as *mut f32,
308                n as i32,
309            ))?;
310            let (mut iou, mut v, mut alpha) = (vec![0.0f32; n], vec![0.0f32; n], vec![0.0f32; n]);
311            iou_buf.to_host(&mut iou)?;
312            v_buf.to_host(&mut v)?;
313            alpha_buf.to_host(&mut alpha)?;
314            Ok((iou, v, alpha))
315        }
316
317        #[allow(clippy::too_many_arguments)]
318        fn ciou_bwd_gpu(
319            &self, device: &CudaDevice<'_>,
320            prog: &teeny_cuda::device::program::CudaProgram<'_, YoloCiouLossBackward<f32>>,
321            dy: &[f32], pred: &[f32], target: &[f32],
322            iou: &[f32], v: &[f32], alpha: &[f32], n: usize,
323        ) -> Result<Vec<f32>> {
324            let mut dy_buf     = device.buffer::<f32>(n)?;
325            let mut pred_buf   = device.buffer::<f32>(4 * n)?;
326            let mut target_buf = device.buffer::<f32>(4 * n)?;
327            let mut iou_buf    = device.buffer::<f32>(n)?;
328            let mut v_buf      = device.buffer::<f32>(n)?;
329            let mut alpha_buf  = device.buffer::<f32>(n)?;
330            let d_pred_buf     = device.buffer::<f32>(4 * n)?;
331            dy_buf.to_device(dy)?; pred_buf.to_device(pred)?; target_buf.to_device(target)?;
332            iou_buf.to_device(iou)?; v_buf.to_device(v)?; alpha_buf.to_device(alpha)?;
333            let cfg = testing::launch_config_with_grid(n.div_ceil(self.block_n as usize), prog);
334            device.launch(prog, &cfg, (
335                dy_buf.as_device_ptr()     as *mut f32,
336                pred_buf.as_device_ptr()   as *mut f32,
337                target_buf.as_device_ptr() as *mut f32,
338                iou_buf.as_device_ptr()    as *mut f32,
339                v_buf.as_device_ptr()      as *mut f32,
340                alpha_buf.as_device_ptr()  as *mut f32,
341                d_pred_buf.as_device_ptr() as *mut f32,
342                n as i32,
343            ))?;
344            let mut out = vec![0.0f32; 4 * n];
345            d_pred_buf.to_host(&mut out)?;
346            Ok(out)
347        }
348
349        #[allow(clippy::too_many_arguments)]
350        fn cls_bwd_gpu(
351            &self, device: &CudaDevice<'_>,
352            prog: &teeny_cuda::device::program::CudaProgram<'_, YoloBceClsLossBackward<f32>>,
353            dy: &[f32], pred: &[f32], target: &[f32],
354            n: usize, c: usize,
355        ) -> Result<Vec<f32>> {
356            let mut dy_buf     = device.buffer::<f32>(n)?;
357            let mut pred_buf   = device.buffer::<f32>(c * n)?;
358            let mut target_buf = device.buffer::<f32>(c * n)?;
359            let d_pred_buf     = device.buffer::<f32>(c * n)?;
360            dy_buf.to_device(dy)?; pred_buf.to_device(pred)?; target_buf.to_device(target)?;
361            let cfg = testing::launch_config_with_grid(n.div_ceil(self.block_n as usize), prog);
362            device.launch(prog, &cfg, (
363                dy_buf.as_device_ptr()     as *mut f32,
364                pred_buf.as_device_ptr()   as *mut f32,
365                target_buf.as_device_ptr() as *mut f32,
366                d_pred_buf.as_device_ptr() as *mut f32,
367                n as i32,
368                c as i32,
369            ))?;
370            let mut out = vec![0.0f32; c * n];
371            d_pred_buf.to_host(&mut out)?;
372            Ok(out)
373        }
374    }
375}
376
377#[cfg(feature = "cuda")]
378pub use cuda_impl::Yolo26Loss;
vision_rs/models/yolo/loss/yolo26.rs

vision_rs/models/yolo/loss/
yolo26.rs
