ultralytics-rust

Ultralytics YOLO 全系列的纯 Rust 复刻：训练、验证、预测、导出全链路，零 Python 依赖 （唯一的深度学习底座是 libtorch 的 Rust 绑定 tch，与 PyTorch 官方同一 C++ 运行时）。

# 加载官方 yolov8n.pt（纯 Rust pickle 解析器，无需 Python/官方库）直接推理
yolo detect predict model=yolov8n.pt source=coco8/images/val save=true

# COCO 格式数据训练（SGD/AdamW、warmup+线性/余弦调度、EMA、mosaic 全套增强）
yolo detect train data=coco8.yaml model=yolov8n.pt epochs=100 imgsz=640
# 增强超参全套透传（未指定的键取官方 default.yaml 默认值）
yolo detect train data=coco8.yaml model=yolov8n.pt mosaic=0.0 mixup=0.2 fliplr=0.3 hsv_v=0.5

# 验证（COCO mAP 10 档 + P/R，对拍官方 ap_per_class 语义）
yolo detect val data=coco8.yaml model=yolov8n.pt

# 导出：safetensors 权重包 / 纯 Rust 生成的 ONNX（图内解码对齐官方默认导出）
yolo detect export model=yolov8n.pt format=onnx imgsz=640

作为 crate 依赖（在你的项目里训练/推理）

[dependencies]
ultralytics-rust = "0.1"

use ultralytics_rust::prelude::*;

fn main() -> anyhow::Result<()> {
    // 官方权重直接加载（48 个模型 YAML 已内嵌，开箱即用）
    let mut model = YoloModel::build("yolo26n", Some(80), Device::Cpu, true)?;

    // 训练
    let cfg = TrainCfg { epochs: 100, imgsz: 640, ..Default::default() };
    let report = ul_engine::train(&mut model, std::path::Path::new("coco8.yaml"), &cfg)?;
    println!("mAP50-95 = {:.4}", report.final_metrics.map50_95);

    // 推理
    let dets = ul_engine::predict(&model, &[std::path::PathBuf::from("image.jpg")], 640, Device::Cpu, 0.25, None)?;
    for o in &dets {
        for d in &o.dets {
            println!("cls={} conf={:.2} box={:?}", d.cls, d.score, d.xyxy);
        }
    }
    Ok(())
}

也可只用子 crate（ultralytics-rust-nn 纯架构/权重加载、ultralytics-rust-engine 训练引擎等）， 按需取用。

支持范围

模型家族（官方 YAML 驱动，ultralytics/cfg/models/** 直接嵌入二进制，119 个文件）：

48/48 个官方 YAML 全部构建+前向通过（cargo test -p ul-nn --test parse_all）。 延期（roadmap）：RT-DETR（transformer 栈）、YOLO-World/YOLOE（CLIP 文本编码）、 yolo26-depth（TorchVision 桥）、SAVPE。

里程碑验证（全部可复现，见 crates/*/tests）

与官方的能力对齐

• parse_model：[from, repeats, module, args] 逐行同构；width/depth 缩放 （make_divisible ceil×8）、C3k2 的 m/l/x c3k 特判、SPPF legacy 补 SiLU、 ResNetLayer 的 [c1,c2,s,is_first,n] 参数序
• 模块动物园：Conv/DWConv/Conv2/GhostConv/LightConv/ConvTranspose/Focus、 C1/C2/C2f/C3/C3x/C3f/C3Ghost/C3k/C3k2、Bottleneck/GhostBottleneck/BottleneckCSP、 SPP/SPPF、RepC3/RepConv/RepNCSPELAN4/ELAN1、AConv/ADown/SPPELAN、CBLinear/CBFuse、 SCDown/CIB/C2fCIB、RepVGGDW、PSA 全家（Attention/PSABlock/PSA/C2PSA/C2fPSA）、 A2C2f（YOLOv12 Area-Attention）、HGStem/HGBlock、ResNetLayer、MaxPool2d/ZeroPad2d
• 五任务头：Detect（含 one2one 双头）、Segment(+Proto26)、Pose(+Pose26 的 cv4_kpts/cv4_sigma 结构)、OBB(+OBB26 原始角度)、v10Detect、Classify、SemanticSegment
• 损失：v8DetectionLoss（TAL + BCE + 成对 CIoU + DFL 量化 CE）、E2E 双头 （o2m topk10 + o2o topk1，对齐 E2EDetectLoss）、reg_max=1 的归一化 L1 路径 （官方 use_dfl=False 分支）、分类 CE+label smoothing； 权重 box=7.5/cls=0.5/dfl=1.5 与官方一致；CIoU 为官方语义 （rho²=(dx²+dy²)/4、c²=对角线平方 cw²+ch²）；BCE 用 softplus 稳定形式
• TAL 分配器：官方 8.4 语义——gt 主序 (bs,m,A)、topk scatter 计数、 一锚多 gt 保留最大 IoU、小框放大 stride_val=stride[1]、norm align parity、 CIoU clamp≥0 后取 align（负 CIoU 在 ^β 下会翻正）
• 数据管线：mosaic4 / mixup / 随机透视（translate/scale/shear/degrees/perspective）/ HSV（真 HSV 转换）/ fliplr（COCO17 关键点交换表）/ flipud / BGR / close_mosaic / letterbox（114 灰边）；以上 14 项超参全部经 CLI 透传
• 评测：COCO 10 档 IoU + trapz 101 点插值 + smooth-F1 最优点 P/R（ap_per_class 忠实移植）
• 训练器：官方三参数组（bias 无 wd+warmup_bias_lr=0.1 起步 / norm 无 wd / weight 有 wd×batch÷nbs）手写 SGD+AdamW（含 AdamW 解耦 wd、bias correction）、 nbs 梯度累积（每 max(1,⌈nbs/batch⌉) 步 clip+step）、multi-scale（[0.5,1.5]×imgsz 32 对齐）、 auto 优化器（total_iters≤10000 → AdamW+lr_fit，否则 SGD）、 warmup（lr 0→lr0、momentum 0.8→目标）、线性/余弦调度、裁剪 grad=10、 EMA 逐迭代 ramp（0.9999·(1−e^(−i/2000))）、close_mosaic 末段关闭、 fitness 最优检查点 + patience 早停 + metrics.jsonl（box/cls/dfl 等分量）
• 任务损失：v8SegmentationLoss（proto 双线性上采样 + einsum 掩码 + bbox 裁剪 BCE/面积 ÷fg 数，×box gain）、v8PoseLoss（格单位 kpt 解码 + KeypointLoss((2σ)²·area·2) + kobj BCE，×pose=12/kobj=1）、v8OBBLoss （RotatedTaskAlignedAssigner 角点投影 + probiou(floor=0.01) + rbox2dist DFL + sin²(2Δθ)·宽高比权重角度损失，×angle=1）
• 任务评测：seg=掩码 IoU 矩阵 AP（proto@coef 上采样 + bbox 裁剪）、 pose=OKS 矩阵 AP（COCO17 σ）、obb=probiou 矩阵 AP；均走官方 match_predictions（10 档 IoU）+ ap_per_class

权重入口

• 官方 .pt 直接加载（pickle VM + zip storage 重建 + 按名对位 VarStore）
• 训练产物：weights.safetensors + meta.json（task/nc/names/stride/yaml_file）， 可用 load_checkpoint 复原继续训练/推理
• 导出：format=safetensors 或 format=onnx（解码后输出，语义同官方默认导出）

Crate 结构

crates/
├── ul-core         纯逻辑：几何/框运算、模型+数据集 YAML 解析（官方 cfg build.rs 嵌入）、
│                   COCO mAP（trapz 对拍 numpy）、纯 CPU NMS
├── ul-nn           张量层：模块动物园（参数名与 ultralytics state_dict 逐字对齐）、
│                   parse_model、官方 .pt 加载（pickle VM）、TAL、ONNX 导出
├── ul-data         图像 IO、YOLO 数据集（detect/pose）、增强管线
├── ul-engine       v8/E2E 损失、trainer（auto 优化器/EMA/调度/裁剪）、validator、predictor
└── ultralytics-rs  `yolo` CLI（train/predict/val/export，key=value 语法）

构建与运行

# 需 libtorch 2.11（与 tch 0.24 硬配对）：设 LIBTORCH 指向本机 libtorch
export LIBTORCH=/path/to/libtorch
cargo build --release

# 官方权重（首次自动下载到 weights/）
./target/release/yolo detect predict model=yolov8n.pt source=<图片或目录> save=true

CUDA：libtorch CUDA 版即可（tch 运行时探测），无卡自动回退 CPU。

已知边界（诚实清单）

• 训练侧未实现：AMP（fp16 autocast——tch 绑定未暴露 autocast 上下文）、 多卡 DDP（tch 未暴露 c10d）；rect 推理已实现（predict 32 对齐矩形 letterbox），rect 训练（宽高比分批）未实现
• 增强：copy_paste（flip 模式，IoA<0.3 过滤 + 多边形掩码贴像素）、cutmix、 erasing（classify 管线）已实现并透传（CLI copy_paste=/cutmix=/erasing=）； copy_paste 的 mixup 模式未实现
• 任务面：detect/classify/segment/pose/obb 训练+评测闭环全通；三任务损失 全部分量与官方 8.4.157 逐位/精确对齐（seg/pose/obb 的 detect 分量 + seg/kobj/ angle/pose 附加分量）；官方 pt 推理加载支持 ckpt nc 覆盖（obb=DOTA 15 等）； classify 训练支持目录布局数据集（root//，类按目录名排序）； yolo26 训练语义未与官方曲线对拍
• ONNX 导出：v5-v11 标准 Detect（图内解码）+ v10/yolo26 e2e（图内两级 topk→(1,300,6)）+ PSA 注意力族 + seg/pose/obb 头全图内解码 （seg (1,116,A)、pose (1,56,A) 含 kpt 图内 (y·2+anchor−0.5)·stride 解码+v=σ(v)、 obb (1,20,A) 含 dist2rbox 旋转解码+(sigmoid−0.25)π 角度；均对拍官方 ≤1e-3，checker+ORT 通过）；A2C2f（v12）显式报错（roadmap）
• 数据集缓存：labels 解析结果持久化 cache.json（构造时全量解析一次， 二次启动零解析；条目数不匹配自动失效重建）
• classify 训练：支持目录布局数据集（data=，子目录即类）；COCO 风格 txt 标签不适用于 classify（官方同样分目录）
• 曲线级对拍：与官方 trainer 的逐 epoch 收敛对比脚本见 scripts/curve_parity.py（同参数同数据双方训练），尚未跑 COCO 全量基准
• pickle VM：支持 torch.save 协议 2-4 的 torch 标准子集； .pt 若含非常规 opcode 会显式报错（不支持静默错位）

许可

AGPL-3.0（与 Ultralytics 上游一致——本项目为 Ultralytics 的衍生实现）。 官方 cfg YAML 与模型设计版权归 Ultralytics。