[{"data":1,"prerenderedAt":36},["ShallowReactive",2],{"chapter:vision-rs\u002Fcore-concepts\u002Ftraining.json":3},{"project":4,"route":5,"title":6,"titleHtml":6,"navTitle":6,"part":7,"sourcePath":8,"editUrl":9,"html":10,"toc":11,"hasMermaid":29,"prev":30,"next":33},"vision-rs","\u002Fvision-rs\u002Fcore-concepts\u002Ftraining","Training","Core Concepts","core-concepts\u002Ftraining.md","https:\u002F\u002Fgithub.com\u002Fteenygrad\u002Fvision-rs\u002Fedit\u002Fmain\u002Fbook\u002Fsrc\u002Fcore-concepts\u002Ftraining.md","\u003Cp>Training support lives behind the \u003Ccode>training\u003C\u002Fcode> feature (on by default) and,\nfor the GPU loss kernels, the \u003Ccode>cuda\u003C\u002Fcode> feature. \u003Ccode>vision_rs::models::yolo::loss\u003C\u002Fcode>\nimplements target assignment and the loss functions needed to train YOLO26\nfrom raw model outputs against ground-truth boxes.\u003C\u002Fp>\n\u003Ch2 id=\"anchor-grid\">Anchor grid\u003C\u002Fh2>\n\u003Cp>\u003Ccode>AnchorGrid\u003C\u002Fcode> (in \u003Ccode>loss::anchor\u003C\u002Fcode>) precomputes anchor centres for all three\nFPN levels (strides 8\u002F16\u002F32), flattened into one array ordered level-by-level\nthen row-major within each level. \u003Ccode>AnchorGrid::yolo26(img_h, img_w)\u003C\u002Fcode> builds\nthe grid for a given input resolution.\u003C\u002Fp>\n\u003Ch2 id=\"target-assignment\">Target assignment\u003C\u002Fh2>\n\u003Cp>\u003Ccode>TaskAlignedAssigner\u003C\u002Fcode> (in \u003Ccode>loss::assign\u003C\u002Fcode>) is a simplified, CPU-side\nimplementation of ultralytics’ TaskAlignedAssigner. For each ground-truth\nbox, it scores every anchor as:\u003C\u002Fp>\n\u003Cpre class=\"code-panel\" data-lang=\"text\">\u003Ccode>score = cls_score^alpha * iou^beta\n\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>and assigns the top-\u003Ccode>k\u003C\u002Fcode> anchors per GT as positives (conflicts — multiple\nGTs claiming the same anchor — are broken by highest score). The result\n(\u003Ccode>AssignResult\u003C\u002Fcode>) carries, per anchor: whether it’s positive, the assigned\nGT box\u002Fclass, and a \u003Cem>soft target\u003C\u002Fem> — \u003Ccode>(align \u002F max_align_for_gt) * max_iou_for_gt\u003C\u002Fcode> — used as both the soft classification label and the box\nloss weight, matching ultralytics’ E2ELoss normalisation.\u003C\u002Fp>\n\u003Ch2 id=\"yolo26loss-cuda\">\u003Ccode>Yolo26Loss\u003C\u002Fcode> (CUDA)\u003C\u002Fh2>\n\u003Cp>\u003Ccode>Yolo26Loss::new(img_h, img_w, nc, cap)\u003C\u002Fcode> builds the loss state: the anchor\ngrid, a default assigner (\u003Ccode>top_k = 10\u003C\u002Fcode>) for the one2many head, and a\n\u003Ccode>top_k = 1\u003C\u002Fcode> assigner for the one2one head.\u003C\u002Fp>\n\u003Cpre data-lang=\"rust\" class=\"shiki teeny-datasheet\" style=\"background-color:#16181a;color:#e6e8e3\" tabindex=\"0\">\u003Ccode>\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">pub\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> fn\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\"> compute_grads\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">    &\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\">self\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> device\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">CudaDevice\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><'\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">_\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    boxes\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">],\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> scores\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">],\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    gt_boxes_b\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">;\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> 4\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">]>],\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> gt_cls_b\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">usize\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>],\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">)\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> ->\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> anyhow\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Result\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><(\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>)>;\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003C\u002Fspan>\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>Compiles and runs the CIoU and classification-loss forward\u002Fbackward kernels\n(see \u003Ca href=\"\u002Fvision-rs\u002Fkernels-and-performance\u002Fcustom-kernels\">Custom Kernels\u003C\u002Fa>) for a\nsingle batch, returning \u003Ccode>(d_boxes, d_scores)\u003C\u002Fcode> gradients ready to backprop\ninto the model graph.\u003C\u002Fp>\n\u003Ch3 id=\"dual-head-training\">Dual-head training\u003C\u002Fh3>\n\u003Cpre data-lang=\"rust\" class=\"shiki teeny-datasheet\" style=\"background-color:#16181a;color:#e6e8e3\" tabindex=\"0\">\u003Ccode>\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">pub\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> fn\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\"> compute_grads_dual\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">    &\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\">self\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> device\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">CudaDevice\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><'\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">_\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    boxes_o2m\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">],\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> scores_o2m\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">],\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    boxes_o2o\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">],\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> scores_o2o\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">],\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    gt_boxes_b\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">;\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> 4\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">]>],\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> gt_cls_b\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">usize\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>],\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    w_o2m\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> w_o2o\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">)\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> ->\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> anyhow\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Result\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><(\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Vec\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">f32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>)>;\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003C\u002Fspan>\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>Runs TAL assignment independently for both heads (their own assigners),\nscaling the resulting gradients by \u003Ccode>w_o2m\u003C\u002Fcode>\u002F\u003Ccode>w_o2o\u003C\u002Fcode>. Pair this with\n\u003Ca href=\"\u002Fvision-rs\u002Fcore-concepts\u002Fyolo26-architecture\">\u003Ccode>yolo26_dual\u003C\u002Fcode>\u003C\u002Fa> for the matching dual-head forward\npass.\u003C\u002Fp>\n\u003Cp>\u003Cstrong>Loss weight schedule\u003C\u002Fstrong> (ultralytics-style): \u003Ccode>w_o2m = 1.0\u003C\u002Fcode> constant\nthroughout training; \u003Ccode>w_o2o = step \u002F total_steps\u003C\u002Fcode>, ramping 0→1 linearly so\nthe one2one head — the one actually used at inference — gradually takes\nover by the end of training. The caller controls the schedule; \u003Ccode>Yolo26Loss\u003C\u002Fcode>\njust applies the weights you pass in.\u003C\u002Fp>\n\u003Cblockquote>\n\u003Cp>One2many is traced \u003Cem>before\u003C\u002Fem> one2one in \u003Ccode>yolo26_dual\u003C\u002Fcode>’s forward closure —\nthe training loop relies on this ordering for stable DAG node\nidentification. If you’re writing a custom training loop against the\ntraced graph directly, don’t reorder the two head calls.\u003C\u002Fp>\n\u003C\u002Fblockquote>\n\u003Ch2 id=\"the-yolo26-examples-traindebugtrain-subcommands\">The \u003Ccode>yolo26\u003C\u002Fcode> example’s \u003Ccode>Train\u003C\u002Fcode>\u002F\u003Ccode>DebugTrain\u003C\u002Fcode> subcommands\u003C\u002Fh2>\n\u003Cp>\u003Ccode>examples\u002Fyolo26.rs\u003C\u002Fcode> has a full CLI training loop wired up against this\nAPI — see its \u003Ccode>Train\u003C\u002Fcode> subcommand for a working reference implementation,\nand \u003Ccode>DebugTrain\u003C\u002Fcode> for a variant that dumps intermediate gradient statistics\n(useful when debugging a new loss\u002Fkernel change).\u003C\u002Fp>\n",[12,16,19,22,26],{"id":13,"text":14,"level":15},"anchor-grid","Anchor grid",2,{"id":17,"text":18,"level":15},"target-assignment","Target assignment",{"id":20,"text":21,"level":15},"yolo26loss-cuda","Yolo26Loss (CUDA)",{"id":23,"text":24,"level":25},"dual-head-training","Dual-head training",3,{"id":27,"text":28,"level":15},"the-yolo26-examples-traindebugtrain-subcommands","The yolo26 example’s Train\u002FDebugTrain subcommands",false,{"title":31,"titleHtml":31,"route":32},"The YOLO26 Model","\u002Fvision-rs\u002Fcore-concepts\u002Fyolo26-architecture",{"title":34,"titleHtml":34,"route":35},"Custom Kernels","\u002Fvision-rs\u002Fkernels-and-performance\u002Fcustom-kernels",1786271830488]