diff --git a/examples/YOLO-Master-Edge-Deployment/README.md b/examples/YOLO-Master-Edge-Deployment/README.md index 435b4777..036a7183 100644 --- a/examples/YOLO-Master-Edge-Deployment/README.md +++ b/examples/YOLO-Master-Edge-Deployment/README.md @@ -77,6 +77,7 @@ examples/YOLO-Master-Edge-Deployment/build-ort/yolo_master_edge_benchmark \ --images /path/to/VisDrone/images/val \ --profile visdrone \ --imgsz 960 \ + --threads 4 \ --limit 500 \ --output benchmark_onnx.csv ``` @@ -107,6 +108,7 @@ examples/YOLO-Master-Edge-Deployment/build-ncnn/yolo_master_edge_benchmark \ --images /path/to/VisDrone/images/val \ --profile visdrone \ --imgsz 960 \ + --threads 4 \ --limit 500 \ --output benchmark_ncnn.csv ``` @@ -137,11 +139,17 @@ examples/YOLO-Master-Edge-Deployment/build-mnn/yolo_master_edge_benchmark \ --images /path/to/VisDrone/images/val \ --profile visdrone \ --imgsz 960 \ + --threads 4 \ --limit 500 \ --output benchmark_mnn.csv ``` `--images` accepts either a directory of images or a text file with one image path per line. +`--threads` configures the CPU worker count for ONNX Runtime, NCNN, and MNN. Keep it identical across backends for a +fair CPU comparison. Thread configuration is applied before each runtime loads its model. + +For MNN, the input session is resized only when the input tensor shape changes. Fixed-shape benchmark loops therefore +exclude repeated session-resize overhead. ## Benchmark CSV Output diff --git a/examples/YOLO-Master-Edge-Deployment/cpp/backends/backend.h b/examples/YOLO-Master-Edge-Deployment/cpp/backends/backend.h index be7db521..017ef109 100644 --- a/examples/YOLO-Master-Edge-Deployment/cpp/backends/backend.h +++ b/examples/YOLO-Master-Edge-Deployment/cpp/backends/backend.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -11,6 +12,7 @@ struct Tensor { class Backend { public: virtual ~Backend() = default; + virtual void set_num_threads(int threads) = 0; virtual void load(const std::string& model_path) = 0; virtual Tensor infer(const Tensor& input) = 0; virtual std::string name() const = 0; diff --git a/examples/YOLO-Master-Edge-Deployment/cpp/backends/mnn_backend.cpp b/examples/YOLO-Master-Edge-Deployment/cpp/backends/mnn_backend.cpp index 0ba65f96..a7c3bcd7 100644 --- a/examples/YOLO-Master-Edge-Deployment/cpp/backends/mnn_backend.cpp +++ b/examples/YOLO-Master-Edge-Deployment/cpp/backends/mnn_backend.cpp @@ -37,6 +37,13 @@ void validate_input(const Tensor& input) { MnnBackend::MnnBackend() = default; +void MnnBackend::set_num_threads(int threads) { + if (threads <= 0) { + throw std::invalid_argument("MNN thread count must be positive"); + } + num_threads_ = threads; +} + MnnBackend::~MnnBackend() { #ifdef WITH_MNN if (interpreter_) { @@ -60,7 +67,7 @@ void MnnBackend::load(const std::string& model_path) { MNN::ScheduleConfig config; config.type = MNN_FORWARD_CPU; - config.numThread = 1; + config.numThread = num_threads_; session_ = interpreter_->createSession(config); if (!session_) { throw std::runtime_error("failed to create MNN session: " + model_path_); @@ -87,8 +94,11 @@ Tensor MnnBackend::infer(const Tensor& input) { } std::vector dims(input.shape.begin(), input.shape.end()); - interpreter_->resizeTensor(input_tensor_, dims); - interpreter_->resizeSession(session_); + if (input_shape_ != input.shape) { + interpreter_->resizeTensor(input_tensor_, dims); + interpreter_->resizeSession(session_); + input_shape_ = input.shape; + } auto* tmp_input = MNN::Tensor::create( dims, diff --git a/examples/YOLO-Master-Edge-Deployment/cpp/backends/mnn_backend.h b/examples/YOLO-Master-Edge-Deployment/cpp/backends/mnn_backend.h index 5d297ec9..d0172d01 100644 --- a/examples/YOLO-Master-Edge-Deployment/cpp/backends/mnn_backend.h +++ b/examples/YOLO-Master-Edge-Deployment/cpp/backends/mnn_backend.h @@ -10,12 +10,15 @@ class MnnBackend final : public Backend { public: MnnBackend(); ~MnnBackend() override; + void set_num_threads(int threads) override; void load(const std::string& model_path) override; Tensor infer(const Tensor& input) override; std::string name() const override; private: std::string model_path_; + int num_threads_ = 4; + std::vector input_shape_; #ifdef WITH_MNN MNN::Interpreter* interpreter_ = nullptr; MNN::Session* session_ = nullptr; diff --git a/examples/YOLO-Master-Edge-Deployment/cpp/backends/ncnn_backend.cpp b/examples/YOLO-Master-Edge-Deployment/cpp/backends/ncnn_backend.cpp index 2c678d23..c34958ae 100644 --- a/examples/YOLO-Master-Edge-Deployment/cpp/backends/ncnn_backend.cpp +++ b/examples/YOLO-Master-Edge-Deployment/cpp/backends/ncnn_backend.cpp @@ -163,6 +163,13 @@ NcnnBackend::NcnnBackend() = default; NcnnBackend::~NcnnBackend() = default; +void NcnnBackend::set_num_threads(int threads) { + if (threads <= 0) { + throw std::invalid_argument("NCNN thread count must be positive"); + } + num_threads_ = threads; +} + void NcnnBackend::load(const std::string& model_path) { if (model_path.empty()) { throw std::invalid_argument("NCNN model path is empty"); @@ -176,7 +183,7 @@ void NcnnBackend::load(const std::string& model_path) { net_.reset(new ncnn::Net()); net_->opt.use_vulkan_compute = false; - net_->opt.num_threads = 1; + net_->opt.num_threads = num_threads_; if (net_->load_param(param_path_.c_str()) != 0) { throw std::runtime_error("failed to load NCNN param file: " + param_path_); diff --git a/examples/YOLO-Master-Edge-Deployment/cpp/backends/ncnn_backend.h b/examples/YOLO-Master-Edge-Deployment/cpp/backends/ncnn_backend.h index b70dd94b..8680e29f 100644 --- a/examples/YOLO-Master-Edge-Deployment/cpp/backends/ncnn_backend.h +++ b/examples/YOLO-Master-Edge-Deployment/cpp/backends/ncnn_backend.h @@ -11,12 +11,14 @@ class NcnnBackend final : public Backend { public: NcnnBackend(); ~NcnnBackend() override; + void set_num_threads(int threads) override; void load(const std::string& model_path) override; Tensor infer(const Tensor& input) override; std::string name() const override; private: std::string model_path_; + int num_threads_ = 4; #ifdef WITH_NCNN std::string param_path_; std::string bin_path_; diff --git a/examples/YOLO-Master-Edge-Deployment/cpp/backends/onnx_backend.cpp b/examples/YOLO-Master-Edge-Deployment/cpp/backends/onnx_backend.cpp index 731dd2d8..7b8f4438 100644 --- a/examples/YOLO-Master-Edge-Deployment/cpp/backends/onnx_backend.cpp +++ b/examples/YOLO-Master-Edge-Deployment/cpp/backends/onnx_backend.cpp @@ -11,6 +11,13 @@ OnnxBackend::OnnxBackend() { } +void OnnxBackend::set_num_threads(int threads) { + if (threads <= 0) { + throw std::invalid_argument("ONNX Runtime thread count must be positive"); + } + num_threads_ = threads; +} + void OnnxBackend::load(const std::string& model_path) { if (model_path.empty()) { throw std::invalid_argument("ONNX model path is empty"); @@ -18,7 +25,9 @@ void OnnxBackend::load(const std::string& model_path) { model_path_ = model_path; #ifdef WITH_ONNXRUNTIME - session_options_.SetIntraOpNumThreads(1); + session_options_.SetIntraOpNumThreads(num_threads_); + session_options_.SetInterOpNumThreads(1); + session_options_.SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL); session_options_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); session_.reset(new Ort::Session(env_, model_path.c_str(), session_options_)); diff --git a/examples/YOLO-Master-Edge-Deployment/cpp/backends/onnx_backend.h b/examples/YOLO-Master-Edge-Deployment/cpp/backends/onnx_backend.h index 5ad59f2e..045aea54 100644 --- a/examples/YOLO-Master-Edge-Deployment/cpp/backends/onnx_backend.h +++ b/examples/YOLO-Master-Edge-Deployment/cpp/backends/onnx_backend.h @@ -10,12 +10,14 @@ class OnnxBackend final : public Backend { public: OnnxBackend(); + void set_num_threads(int threads) override; void load(const std::string& model_path) override; Tensor infer(const Tensor& input) override; std::string name() const override; private: std::string model_path_; + int num_threads_ = 4; #ifdef WITH_ONNXRUNTIME Ort::Env env_; Ort::SessionOptions session_options_; diff --git a/examples/YOLO-Master-Edge-Deployment/cpp/edge_benchmark.cpp b/examples/YOLO-Master-Edge-Deployment/cpp/edge_benchmark.cpp index 783c05ba..08087c4e 100644 --- a/examples/YOLO-Master-Edge-Deployment/cpp/edge_benchmark.cpp +++ b/examples/YOLO-Master-Edge-Deployment/cpp/edge_benchmark.cpp @@ -28,6 +28,7 @@ struct Args { int warmup = 5; int runs = 1; int limit = 0; + int threads = 4; }; struct TimingRow { @@ -52,6 +53,7 @@ static void print_usage(const char* program) { << "[--warmup 5] " << "[--runs 1] " << "[--limit 500] " + << "[--threads 4] " << "[--output benchmark.csv]\n"; } @@ -100,6 +102,8 @@ static Args parse_args(int argc, char** argv) { args.runs = std::stoi(value); } else if (key == "--limit") { args.limit = std::stoi(value); + } else if (key == "--threads") { + args.threads = std::stoi(value); } else { std::cerr << "Unknown argument: " << key << "\n"; print_usage(argv[0]); @@ -120,7 +124,8 @@ static Args parse_args(int argc, char** argv) { std::cerr << "Invalid --profile: " << args.profile << "\n"; std::exit(2); } - if (args.imgsz <= 0 || args.warmup < 0 || args.runs <= 0 || args.limit < 0) { + if (args.imgsz <= 0 || args.warmup < 0 || args.runs <= 0 || args.limit < 0 || + args.threads <= 0) { std::cerr << "Invalid numeric argument\n"; std::exit(2); } @@ -251,6 +256,7 @@ int main(int argc, char** argv) { const Args args = parse_args(argc, argv); const auto images = collect_images(args.images, args.limit); auto backend = create_backend(args.backend); + backend->set_num_threads(args.threads); backend->load(args.model); const Tensor warmup_input = preprocess_image(images.front(), args.imgsz, args.imgsz).input; @@ -295,6 +301,7 @@ int main(int argc, char** argv) { << " model=" << args.model << " profile=" << args.profile << " imgsz=" << args.imgsz + << " threads=" << args.threads << " conf=" << args.conf << " iou=" << args.iou << " output=" << args.output << "\n"; diff --git a/examples/lora_examples/yolo_master_lora_peft_ema_report.md b/examples/lora_examples/yolo_master_lora_peft_ema_report.md new file mode 100644 index 00000000..b0b4a573 --- /dev/null +++ b/examples/lora_examples/yolo_master_lora_peft_ema_report.md @@ -0,0 +1,62 @@ +# YOLO-Master PEFT LoRA EMA 同步实验报告 + +本报告记录 `peft_ema_sync_rtx4060_v1` 协议。该协议用于验证 PEFT LoRA 的非张量 `scaling` +状态同步到 EMA 后，在 Brain Tumor 和 VisDrone 垂类场景中的 rank 扫描结果。 + +六组正式实验基于仓库提交 `a510883` 加本地 PEFT EMA 修复运行；实验完成后，修复提交才重放到 +更新后的 `upstream/main`。因此结果应以本报告列出的完整协议为准，不能套用后续默认配置解释。 + +## 问题与修复 + +启用 `lora_alpha_warmup` 后，在线模型的 LoRA `scaling` 会随 epoch 增长，但 PEFT 0.19.1 将该状态 +保存在普通 Python 字典中，而不是 `state_dict` 张量。标准 EMA 更新因此不会复制它，导致在线模型使用 +LoRA、EMA 验证模型却保持零缩放。典型现象是训练 loss 下降，但 mAP 持续下降或归零。 + +修复在以下生命周期边界同步在线模型与 EMA 的 LoRA `scaling`： + +- 每个 epoch 更新 alpha warmup 后； +- 验证前； +- checkpoint 序列化前； +- 断点续训恢复后。 + +## 实验环境与协议 + +- GPU：NVIDIA GeForce RTX 4060 Laptop GPU（8188 MiB） +- Python：3.11.15 +- PyTorch：2.13.0+cu126 +- PEFT：0.19.1 +- 模型：YOLO-Master-EsMoE-N 预训练权重 +- Rank：`r=4,8,16`，保持 `lora_alpha=2*r` +- Backend：配置为 `auto`，实际解析为 `peft` +- AMP：关闭，避免把数值稳定性问题混入 EMA 修复验证 +- Router/gating：不纳入 LoRA 目标模块 + +Brain Tumor 使用全部训练集、`imgsz=640`、`batch=8`、最多 40 epochs、`patience=15`、 +`lora_alpha_warmup=3`。VisDrone 使用 20% 训练集、完整验证集、`imgsz=640`、`batch=2`、 +30 epochs、`lora_alpha_warmup=5` 和多尺度训练。 + +## Rank 扫描结果 + +| 数据集 | Rank | 完成轮数 | 最佳轮次 | mAP50 | mAP50-95 | 可训练参数 | Adapter 参数 | 时间（分钟） | 日志峰值显存 | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| Brain Tumor | 4 | 17 | 2 | 0.40754 | 0.26097 | 409,174 | 64,000 | 10.38 | 3.79 GB | +| Brain Tumor | 8 | 40 | 33 | 0.47810 | **0.34647** | 473,174 | 128,000 | 25.38 | 3.83 GB | +| Brain Tumor | 16 | 17 | 2 | 0.47357 | 0.31845 | 601,174 | 256,000 | 10.62 | 3.84 GB | +| VisDrone | 4 | 30 | 20 | 0.09152 | 0.04601 | 410,734 | 64,000 | 73.14 | 8.68 GB | +| VisDrone | 8 | 30 | 20 | 0.09799 | 0.04926 | 474,734 | 128,000 | 69.06 | 8.69 GB | +| VisDrone | 16 | 30 | 28 | 0.11454 | **0.05797** | 602,734 | 256,000 | 76.62 | 8.72 GB | + +峰值显存来自训练日志的 `GPU_mem` 最大值；不同 CUDA/PyTorch 版本的内存统计口径可能不同。 +完整机器可读结果见 `yolo_master_lora_peft_ema_results.csv`。 + +## 结论 + +- Brain Tumor 推荐 `r=8`：mAP50-95 最高，且比 `r=16` 少 128,000 个 Adapter 参数。 +- VisDrone 推荐 `r=16`：密集小目标场景从更大的 LoRA 容量中获得了持续收益。 +- 两个场景不存在统一最佳 rank，rank 应根据领域复杂度分别选择。 +- `best.pt` 重新验证结果与训练记录一致，修复后未再出现 LoRA EMA 缩放为零导致的指标崩溃。 + +## 可比性限制 + +本协议不能与仓库中的历史协议直接合并。历史结果可能使用 fallback 后端、AMP、不同 batch、 +不同图像尺寸或不同数据比例。跨协议数值只能作为背景参考，rank 结论应在同一协议内部比较。 diff --git a/examples/lora_examples/yolo_master_lora_peft_ema_results.csv b/examples/lora_examples/yolo_master_lora_peft_ema_results.csv new file mode 100644 index 00000000..6a063439 --- /dev/null +++ b/examples/lora_examples/yolo_master_lora_peft_ema_results.csv @@ -0,0 +1,7 @@ +protocol_id,dataset,rank,alpha,max_epochs,completed_epochs,fraction,amp,batch,imgsz,effective_backend,alpha_warmup,best_epoch,precision,recall,mAP50,mAP50_95,trainable_params,adapter_params,train_time_min,peak_gpu_memory_gb,status +peft_ema_sync_rtx4060_v1,brain_tumor,4,8,40,17,1.0,False,8,640,peft,3,2,0.43386,0.57434,0.40754,0.26097,409174,64000,10.38,3.79,completed_early_stop +peft_ema_sync_rtx4060_v1,brain_tumor,8,16,40,40,1.0,False,8,640,peft,3,33,0.45512,0.77818,0.47810,0.34647,473174,128000,25.38,3.83,completed +peft_ema_sync_rtx4060_v1,brain_tumor,16,32,40,17,1.0,False,8,640,peft,3,2,0.43767,0.75353,0.47357,0.31845,601174,256000,10.62,3.84,completed_early_stop +peft_ema_sync_rtx4060_v1,visdrone,4,8,30,30,0.2,False,2,640,peft,5,20,0.27431,0.14468,0.09152,0.04601,410734,64000,73.14,8.68,completed +peft_ema_sync_rtx4060_v1,visdrone,8,16,30,30,0.2,False,2,640,peft,5,20,0.24676,0.14905,0.09799,0.04926,474734,128000,69.06,8.69,completed +peft_ema_sync_rtx4060_v1,visdrone,16,32,30,30,0.2,False,2,640,peft,5,28,0.30409,0.15179,0.11454,0.05797,602734,256000,76.62,8.72,completed diff --git a/scripts/aggregate_mot_ablation_seeds.py b/scripts/aggregate_mot_ablation_seeds.py new file mode 100644 index 00000000..04fb6320 --- /dev/null +++ b/scripts/aggregate_mot_ablation_seeds.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +"""Aggregate multi-seed MoE/MoT/MoA ablations with completeness checks. + +The training summaries contain seed-dependent accuracy and stability metrics. A +single optional benchmark CSV supplies architecture-dependent latency, FLOPs, +and parameter counts, which do not need to be remeasured for every seed. +""" + +from __future__ import annotations + +import argparse +import csv +import math +import statistics +from collections import defaultdict +from pathlib import Path + + +METRICS = { + "map50_95": ("metrics/mAP50-95(B)", "mAP50-95"), + "map50": ("metrics/mAP50(B)", "mAP50"), + "final_train_total_loss": ("final_train_total_loss",), +} +PROFILE_FIELDS = ("latency_ms_p50", "latency_ms_p95", "latency_ms_p99", "flops_g", "params_m") + + +def read_csv(path: Path) -> list[dict[str, str]]: + with path.open(newline="", encoding="utf-8-sig") as handle: + return list(csv.DictReader(handle)) + + +def as_float(value: object) -> float | None: + try: + result = float(value) if value not in {None, ""} else None + except (TypeError, ValueError): + return None + return result if result is not None and math.isfinite(result) else None + + +def as_bool(value: object) -> bool: + return str(value).strip().lower() in {"1", "true", "yes", "y"} + + +def seed_sort_key(seed: str) -> tuple[int, int | str]: + return (0, int(seed)) if seed.isdigit() else (1, seed) + + +def mean_std(values: list[float]) -> tuple[float, float]: + return statistics.mean(values), statistics.stdev(values) if len(values) > 1 else 0.0 + + +def load_optional_by_key(path: Path | None) -> dict[str, dict[str, str]]: + if not path: + return {} + if not path.exists(): + raise FileNotFoundError(path) + rows = read_csv(path) + if any(not row.get("key") for row in rows): + raise ValueError(f"profile CSV has a row without key: {path}") + return {row["key"]: row for row in rows} + + +def collect_seed_rows( + root: Path, + expected_seeds: list[str] | None = None, + allow_incomplete: bool = False, +) -> dict[str, list[dict[str, str]]]: + summaries = sorted(root.glob("seed_*/summary.csv"), key=lambda path: seed_sort_key(path.parent.name[5:])) + if not summaries: + raise ValueError(f"no seed_*/summary.csv files found under {root}") + + discovered_seeds = [path.parent.name[5:] for path in summaries] + if expected_seeds is not None: + expected = sorted({str(seed) for seed in expected_seeds}, key=seed_sort_key) + if discovered_seeds != expected: + raise ValueError(f"seed mismatch: expected {expected}, found {discovered_seeds}") + + grouped: dict[str, list[dict[str, str]]] = defaultdict(list) + model_sets: dict[str, set[str]] = {} + for summary in summaries: + seed = summary.parent.name[5:] + rows = read_csv(summary) + keys = [row.get("key", "") for row in rows] + if not rows or any(not key for key in keys): + raise ValueError(f"empty or malformed summary: {summary}") + if len(keys) != len(set(keys)): + raise ValueError(f"duplicate model key in {summary}") + model_sets[seed] = set(keys) + for row in rows: + grouped[row["key"]].append({**row, "seed": seed, "source": str(summary)}) + + if not allow_incomplete: + expected_models = set.union(*model_sets.values()) + incomplete = { + seed: sorted(expected_models - keys) + for seed, keys in model_sets.items() + if keys != expected_models + } + if incomplete: + raise ValueError(f"incomplete model coverage by seed: {incomplete}") + return grouped + + +def aggregate( + root: Path, + latency_csv: Path | None = None, + build_csv: Path | None = None, + *, + baseline_key: str = "v10", + expected_seeds: list[str] | None = None, + allow_incomplete: bool = False, + map_gain_threshold: float = 0.01, + latency_reduction_threshold_pct: float = 10.0, +) -> list[dict[str, object]]: + grouped = collect_seed_rows(root, expected_seeds=expected_seeds, allow_incomplete=allow_incomplete) + latency = load_optional_by_key(latency_csv) + builds = load_optional_by_key(build_csv) + output: list[dict[str, object]] = [] + + for key, rows in sorted(grouped.items()): + ordered_rows = sorted(rows, key=lambda row: seed_sort_key(row["seed"])) + item: dict[str, object] = { + "key": key, + "label": ordered_rows[0].get("label", key), + "n_seeds": len(ordered_rows), + "seeds": ",".join(row["seed"] for row in ordered_rows), + "nan_any": any(as_bool(row.get("nan_detected")) for row in ordered_rows), + "loss_diverged_any": any(as_bool(row.get("loss_diverged")) for row in ordered_rows), + } + for output_name, source_names in METRICS.items(): + values = [] + for row in ordered_rows: + value = next((parsed for name in source_names if (parsed := as_float(row.get(name))) is not None), None) + if value is not None: + values.append(value) + if values: + mean, std = mean_std(values) + item.update( + { + f"{output_name}_mean": mean, + f"{output_name}_std": std, + f"{output_name}_min": min(values), + f"{output_name}_max": max(values), + } + ) + + profile = {**builds.get(key, {}), **latency.get(key, {})} + fallback = ordered_rows[0] + for field in PROFILE_FIELDS: + value = as_float(profile.get(field)) + if value is None: + value = as_float(fallback.get(field)) + if value is not None: + item[field] = value + output.append(item) + + baseline = next((row for row in output if row["key"] == baseline_key), None) + if baseline is None: + raise ValueError(f"baseline key {baseline_key!r} is absent") + baseline_map = as_float(baseline.get("map50_95_mean")) + baseline_p50 = as_float(baseline.get("latency_ms_p50")) + for row in output: + current_map = as_float(row.get("map50_95_mean")) + current_p50 = as_float(row.get("latency_ms_p50")) + map_delta = current_map - baseline_map if current_map is not None and baseline_map is not None else None + latency_delta = ( + (current_p50 - baseline_p50) / baseline_p50 * 100 + if current_p50 is not None and baseline_p50 not in {None, 0} + else None + ) + if map_delta is not None: + row["map50_95_delta_vs_baseline"] = map_delta + if latency_delta is not None: + row["latency_delta_pct_vs_baseline"] = latency_delta + row["meaningful_gain"] = bool( + (map_delta is not None and map_delta > map_gain_threshold) + or (latency_delta is not None and latency_delta < -latency_reduction_threshold_pct) + ) + return output + + +def write_csv(path: Path, rows: list[dict[str, object]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fieldnames = sorted({key for row in rows for key in row}) + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + +def fmt(row: dict[str, object], key: str, digits: int = 4) -> str: + value = as_float(row.get(key)) + return f"{value:.{digits}f}" if value is not None else "N/A" + + +def write_markdown(path: Path, rows: list[dict[str, object]], title: str, note: str | None = None) -> None: + lines = [ + f"# {title}", + "", + "| Model | Seeds | mAP50-95 mean±std | mAP50 mean±std | P50/P95/P99 ms | " + "Actual FLOPs (G) | Params (M) | Stable | Gain gate |", + "|---|---:|---:|---:|---:|---:|---:|:---:|:---:|", + ] + for row in rows: + map95 = f"{fmt(row, 'map50_95_mean')}±{fmt(row, 'map50_95_std')}" + map50 = f"{fmt(row, 'map50_mean')}±{fmt(row, 'map50_std')}" + latency = "/".join(fmt(row, key, 2) for key in ("latency_ms_p50", "latency_ms_p95", "latency_ms_p99")) + stable = "yes" if not row.get("nan_any") and not row.get("loss_diverged_any") else "no" + gain = "pass" if row.get("meaningful_gain") else "fail" + lines.append( + f"| {row.get('label', row['key'])} | {row['n_seeds']} | {map95} | {map50} | {latency} | " + f"{fmt(row, 'flops_g', 3)} | {fmt(row, 'params_m', 3)} | {stable} | {gain} |" + ) + if note: + lines.extend(["", f"> {note}"]) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", type=Path, required=True, help="Directory containing seed_*/summary.csv.") + parser.add_argument("--latency-csv", type=Path, help="One hardware-controlled benchmark CSV shared by all seeds.") + parser.add_argument("--build-csv", type=Path) + parser.add_argument("--baseline", default="v10") + parser.add_argument("--expected-seeds", nargs="+") + parser.add_argument("--allow-incomplete", action="store_true") + parser.add_argument("--map-gain-threshold", type=float, default=0.01) + parser.add_argument("--latency-reduction-threshold-pct", type=float, default=10.0) + parser.add_argument("--title", default="MoE/MoT/MoA multi-seed ablation") + parser.add_argument("--note") + parser.add_argument("--out-csv", type=Path) + parser.add_argument("--out-md", type=Path) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + rows = aggregate( + args.root, + args.latency_csv, + args.build_csv, + baseline_key=args.baseline, + expected_seeds=args.expected_seeds, + allow_incomplete=args.allow_incomplete, + map_gain_threshold=args.map_gain_threshold, + latency_reduction_threshold_pct=args.latency_reduction_threshold_pct, + ) + out_csv = args.out_csv or args.root / "aggregate_multiseed.csv" + out_md = args.out_md or args.root / "aggregate_multiseed.md" + write_csv(out_csv, rows) + write_markdown(out_md, rows, title=args.title, note=args.note) + print(f"wrote {out_csv}") + print(f"wrote {out_md}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_coco128_mot_3seed.sh b/scripts/run_coco128_mot_3seed.sh new file mode 100644 index 00000000..d8f9d379 --- /dev/null +++ b/scripts/run_coco128_mot_3seed.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Reproducible COCO128 pilot for the Issue #54 MoE/MoT/MoA comparison. +# Environment variables may override defaults without editing this file. +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PYTHON="${PYTHON:-python3}" +DATA="${DATA:-$ROOT/ultralytics/cfg/datasets/coco128.yaml}" +PROJECT_ROOT="${PROJECT_ROOT:-$ROOT/runs/mot_ablation_coco128_3seed}" +EPOCHS="${EPOCHS:-100}" +IMGSZ="${IMGSZ:-640}" +BATCH="${BATCH:-16}" +WORKERS="${WORKERS:-0}" +DEVICE="${DEVICE:-0}" +SEEDS="${SEEDS:-42 123 3407}" +MODELS="${MODELS:-v10 v10_mot v10_moa v10_moa_mot}" +BENCHMARK="${BENCHMARK:-1}" +WARMUP="${WARMUP:-100}" +REPS="${REPS:-1000}" + +read -r -a seed_args <<<"$SEEDS" +read -r -a model_args <<<"$MODELS" + +mkdir -p "$PROJECT_ROOT" +for seed in "${seed_args[@]}"; do + echo "[$(date --iso-8601=seconds)] starting seed=$seed models=$MODELS" + "$PYTHON" "$ROOT/scripts/compare_mot_ablation.py" \ + --train \ + --models "${model_args[@]}" \ + --data "$DATA" \ + --epochs "$EPOCHS" \ + --imgsz "$IMGSZ" \ + --batch "$BATCH" \ + --workers "$WORKERS" \ + --device "$DEVICE" \ + --seed "$seed" \ + --project "$PROJECT_ROOT/seed_$seed" \ + --exist-ok + echo "[$(date --iso-8601=seconds)] completed seed=$seed" +done + +aggregate_args=( + --root "$PROJECT_ROOT" + --expected-seeds "${seed_args[@]}" + --title "COCO128 MoE/MoT/MoA 3-seed pilot" + --note "COCO128 is a smoke benchmark whose train and validation images overlap; do not present these metrics as full-COCO generalization results." +) + +if [[ "$BENCHMARK" == "1" ]]; then + profile_dir="$PROJECT_ROOT/profile" + "$PYTHON" "$ROOT/scripts/compare_mot_ablation.py" \ + --benchmark \ + --actual-flops \ + --models "${model_args[@]}" \ + --device "$DEVICE" \ + --imgsz "$IMGSZ" \ + --warmup "$WARMUP" \ + --reps "$REPS" \ + --project "$profile_dir" + aggregate_args+=(--latency-csv "$profile_dir/latency_${DEVICE}_${IMGSZ}.csv") +fi + +"$PYTHON" "$ROOT/scripts/aggregate_mot_ablation_seeds.py" "${aggregate_args[@]}" +echo "[$(date --iso-8601=seconds)] all COCO128 runs and aggregation completed" diff --git a/tests/test_aggregate_mot_ablation_seeds.py b/tests/test_aggregate_mot_ablation_seeds.py new file mode 100644 index 00000000..02a497c4 --- /dev/null +++ b/tests/test_aggregate_mot_ablation_seeds.py @@ -0,0 +1,80 @@ +"""Regression tests for strict multi-seed MoT ablation aggregation.""" + +from pathlib import Path + +import pytest + +from scripts.aggregate_mot_ablation_seeds import aggregate, write_markdown + + +HEADER = "key,label,metrics/mAP50(B),metrics/mAP50-95(B),final_train_total_loss,nan_detected,loss_diverged\n" + + +def write_seed(root: Path, seed: int, rows: list[str]) -> None: + seed_dir = root / f"seed_{seed}" + seed_dir.mkdir(parents=True) + (seed_dir / "summary.csv").write_text(HEADER + "".join(rows), encoding="utf-8") + + +def test_aggregate_combines_seed_metrics_and_one_profile(tmp_path: Path): + root = tmp_path / "runs" + write_seed(root, 42, ["v10,baseline,0.2,0.1,5.0,False,False\n", "v10_mot,mot,0.3,0.2,6.0,False,False\n"]) + write_seed(root, 123, ["v10,baseline,0.4,0.3,5.5,False,False\n", "v10_mot,mot,0.5,0.4,6.5,False,False\n"]) + profile = tmp_path / "latency.csv" + profile.write_text( + "key,latency_ms_p50,latency_ms_p95,latency_ms_p99,flops_g,params_m\n" + "v10,10,11,12,8.5,3.4\n" + "v10_mot,20,21,22,12.2,4.0\n", + encoding="utf-8", + ) + + rows = aggregate(root, profile, expected_seeds=["42", "123"]) + by_key = {row["key"]: row for row in rows} + + assert by_key["v10"]["seeds"] == "42,123" + assert by_key["v10"]["map50_95_mean"] == pytest.approx(0.2) + assert by_key["v10"]["map50_95_std"] == pytest.approx(2**0.5 / 10) + assert by_key["v10_mot"]["latency_ms_p99"] == 22.0 + assert by_key["v10_mot"]["meaningful_gain"] is True + + +def test_aggregate_rejects_missing_seed(tmp_path: Path): + root = tmp_path / "runs" + write_seed(root, 42, ["v10,baseline,0.2,0.1,5.0,False,False\n"]) + + with pytest.raises(ValueError, match="seed mismatch"): + aggregate(root, expected_seeds=["42", "123"]) + + +def test_aggregate_rejects_incomplete_model_coverage(tmp_path: Path): + root = tmp_path / "runs" + write_seed(root, 42, ["v10,baseline,0.2,0.1,5.0,False,False\n", "v10_mot,mot,0.3,0.2,6.0,False,False\n"]) + write_seed(root, 123, ["v10,baseline,0.4,0.3,5.5,False,False\n"]) + + with pytest.raises(ValueError, match="incomplete model coverage"): + aggregate(root) + + +def test_markdown_uses_uncertainty_and_pilot_note(tmp_path: Path): + output = tmp_path / "summary.md" + rows = [ + { + "key": "v10", + "label": "baseline", + "n_seeds": 3, + "map50_95_mean": 0.2, + "map50_95_std": 0.01, + "map50_mean": 0.3, + "map50_std": 0.02, + "nan_any": False, + "loss_diverged_any": False, + "meaningful_gain": False, + } + ] + + write_markdown(output, rows, title="pilot", note="smoke only") + text = output.read_text(encoding="utf-8") + + assert "mean±std" in text + assert "0.2000±0.0100" in text + assert "> smoke only" in text diff --git a/tests/test_lora_selective_ema_lifecycle.py b/tests/test_lora_selective_ema_lifecycle.py index 6162c4c7..a460d565 100644 --- a/tests/test_lora_selective_ema_lifecycle.py +++ b/tests/test_lora_selective_ema_lifecycle.py @@ -44,6 +44,31 @@ def forward(self, x): return self.layer(x) +class TinyPeftLayer(nn.Module): + """Minimal PEFT-like layer whose scaling dictionary is absent from state_dict.""" + + def __init__(self): + super().__init__() + self.lora_A = nn.ModuleDict({"default": nn.Linear(2, 2, bias=False)}) + self.lora_B = nn.ModuleDict({"default": nn.Linear(2, 2, bias=False)}) + self.scaling = {"default": 0.0} + + def forward(self, inputs): + return self.lora_B["default"](self.lora_A["default"](inputs)) * self.scaling["default"] + + +class TinyPeftGraph(nn.Module): + def __init__(self): + super().__init__() + self.layer = TinyPeftLayer() + self.lora_enabled = True + self.lora_backend = "peft" + self.lora_runtime_metadata = {"effective_backend": "peft"} + + def forward(self, inputs): + return self.layer(inputs) + + def _trainer(model: nn.Module, *, warmup: int) -> SimpleNamespace: optimizer = torch.optim.SGD((parameter for parameter in model.parameters() if parameter.requires_grad), lr=0.1) trainer = SimpleNamespace( @@ -110,13 +135,39 @@ def test_resume_restores_scheduled_scaling_online_and_ema(start_epoch: int): assert trainer.ema.ema.layer.scaling == pytest.approx(expected) -def test_non_fallback_backend_does_not_rewrite_ema_treatment(): - model = TinyFallbackGraph(backend="peft") - trainer, controller = _controller(model, warmup=0) - trainer.ema.ema.layer.scaling = 0.0 +def test_peft_backend_syncs_scaling_dictionary_to_ema(): + model = TinyPeftGraph() + trainer = _trainer(model, warmup=0) + controller = AdapterRuntimeController(trainer) + trainer.adapter_controller = controller + trainer.ema = ModelEMA(model) + model.layer.scaling["default"] = 3.0 - assert controller.sync_ema_treatment() == 0 - assert trainer.ema.ema.layer.scaling == 0.0 + assert controller.sync_ema_treatment() == 1 + assert trainer.ema.ema.layer.scaling["default"] == 3.0 + + +def test_validation_syncs_peft_treatment_before_selecting_ema_model(): + model = TinyPeftGraph() + trainer = _trainer(model, warmup=0) + controller = AdapterRuntimeController(trainer) + trainer.adapter_controller = controller + trainer.ema = ModelEMA(model) + model.layer.scaling["default"] = 3.0 + trainer._sync_ema_buffers_for_validation = lambda: None + trainer._state_is_finite = lambda _value: True + trainer.best_fitness = None + trainer.loss = torch.tensor(1.0) + + def validator(runtime_trainer): + assert runtime_trainer.ema.ema.layer.scaling["default"] == 3.0 + return {"fitness": 1.0} + + trainer.validator = validator + metrics, fitness = BaseTrainer.validate(trainer) + + assert metrics == {} + assert fitness == 1.0 def test_validation_syncs_fallback_treatment_before_selecting_ema_model(): @@ -157,3 +208,21 @@ def test_checkpoint_serialization_syncs_fallback_treatment(): assert checkpoint["ema"].layer.scaling == model.layer.scaling assert checkpoint["ema"].layer.use_rslora is True + + +def test_checkpoint_serialization_syncs_peft_treatment(): + model = TinyPeftGraph() + trainer = _trainer(model, warmup=0) + trainer.adapter_controller = AdapterRuntimeController(trainer) + trainer.ema = ModelEMA(model) + model.layer.scaling["default"] = 3.0 + trainer.scaler = SimpleNamespace(state_dict=lambda: {}) + trainer.epoch = trainer.start_epoch = 0 + trainer.best_fitness = trainer.fitness = 0.0 + trainer.metrics = {} + trainer.read_results_csv = lambda: {} + + serialized = TrainingRecoveryController(trainer).serialize_checkpoint(include_online_model=True) + checkpoint = torch.load(io.BytesIO(serialized), map_location="cpu", weights_only=False) + + assert checkpoint["ema"].layer.scaling["default"] == 3.0 diff --git a/tests/test_moe_prune_yaml_sync.py b/tests/test_moe_prune_yaml_sync.py index 1de1be4c..f27c5fc2 100644 --- a/tests/test_moe_prune_yaml_sync.py +++ b/tests/test_moe_prune_yaml_sync.py @@ -1,14 +1,18 @@ # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license -"""Regression test: MoEPruner must keep ``model.yaml`` consistent with the pruned -expert count so a pruned model survives a YAML-based rebuild during retraining +"""Regression tests: MoEPruner must keep ``model.yaml`` consistent with the pruned +ES_MOE blocks so a pruned model survives a YAML-based rebuild during retraining (the prune -> LoRA / full fine-tune recovery workflow). -Without the fix, ``DetectionModel(pruned.yaml)`` rebuilds with ES_MOE's default -expert count and ``intersect_dicts`` silently drops the reduced expert/router -weights, re-inflating the model with randomly initialized experts. -""" -import copy +Two failure modes are covered: +1. Expert *count* re-inflation - rebuild uses ES_MOE's default expert count and + ``intersect_dicts`` drops the reduced expert/router weights. +2. Expert *kernel* mismatch - pruning keeps experts with heterogeneous kernels + (e.g. [5, 9]) but a bare rebuild assigns the defaults [3, 5], so the kept + experts' depthwise weights are dropped on a shape mismatch and re-initialized. +The fix writes the full ES_MOE arg list (count + per-expert kernel sizes) into +``model.yaml`` so both survive the rebuild. +""" import yaml from ultralytics.nn.modules.moe.modules import ES_MOE @@ -22,6 +26,14 @@ def _experts(model): return [m.num_experts for _, m in model.named_modules() if isinstance(m, ES_MOE)] +def _kernels(model): + return [ + [e.conv.depthwise.kernel_size[0] for e in m.experts] + for _, m in model.named_modules() + if isinstance(m, ES_MOE) + ] + + def _build_reduced(num_experts): """Build a model whose ES_MOE blocks each have ``num_experts`` (stand-in for a pruned model).""" d = yaml.safe_load(open(CFG)) @@ -31,6 +43,16 @@ def _build_reduced(num_experts): return DetectionModel(d, ch=3, nc=10, verbose=False) +def _build_with_kernels(num_experts, kernels): + """Stand-in for a pruned model that kept experts with non-default kernel sizes.""" + d = yaml.safe_load(open(CFG)) + for layer in d["backbone"]: + if layer[2] == "ES_MOE": + out_ch = layer[3][0] + layer[3] = [out_ch, num_experts, 8, num_experts, True, 0.4, 15, list(kernels)] + return DetectionModel(d, ch=3, nc=10, verbose=False) + + def test_reinflation_without_sync(): """Reproduce the bug: reduced model + un-synced (full) yaml re-inflates on rebuild.""" reduced = _build_reduced(2) @@ -41,12 +63,24 @@ def test_reinflation_without_sync(): def test_sync_preserves_pruned_experts(): - """The fix: MoEPruner._sync_yaml_num_experts writes the reduced count into yaml, - so a YAML rebuild preserves the pruned architecture.""" + """The fix writes the reduced count into yaml so a rebuild preserves it.""" reduced = _build_reduced(2) reduced.yaml = yaml.safe_load(open(CFG)) # start from the un-synced (buggy) yaml MoEPruner._sync_yaml_num_experts(type("D", (), {})(), reduced) # method only uses self for logging es_args = [layer[3] for layer in reduced.yaml["backbone"] if layer[2] == "ES_MOE"] - assert all(a[-1] == 2 for a in es_args), f"yaml not synced: {es_args}" + assert all(a[1] == 2 for a in es_args), f"num_experts not synced: {es_args}" rebuilt = DetectionModel(reduced.yaml, ch=3, nc=10, verbose=False) assert _experts(rebuilt) == [2, 2, 2, 2], "pruned expert count must survive the rebuild" + + +def test_sync_preserves_expert_kernels(): + """The fix also writes per-expert kernel sizes so heterogeneous kept experts survive.""" + reduced = _build_with_kernels(2, [5, 9]) + before = _kernels(reduced) + assert all(k == [5, 9] for k in before), f"stand-in kernels not built: {before}" + reduced.yaml = yaml.safe_load(open(CFG)) # un-synced yaml would rebuild default [3, 5] + MoEPruner._sync_yaml_num_experts(type("D", (), {})(), reduced) + es_args = [layer[3] for layer in reduced.yaml["backbone"] if layer[2] == "ES_MOE"] + assert all(a[-1] == [5, 9] for a in es_args), f"kernels not synced: {es_args}" + rebuilt = DetectionModel(reduced.yaml, ch=3, nc=10, verbose=False) + assert _kernels(rebuilt) == before, "expert kernel sizes must survive the rebuild" diff --git a/ultralytics/engine/extensions/adapters.py b/ultralytics/engine/extensions/adapters.py index 85a22aa3..33a3693a 100644 --- a/ultralytics/engine/extensions/adapters.py +++ b/ultralytics/engine/extensions/adapters.py @@ -224,18 +224,22 @@ def configure_optimizer(self, optimizer) -> None: self.trainer.lora_ortho_batch_counter = 0 def sync_ema_treatment(self) -> int: - """Copy scheduled fallback scaling from online wrappers to matching EMA wrappers.""" + """Copy scheduled LoRA scaling from online adapters to matching EMA adapters.""" metadata = getattr(self.model, "lora_runtime_metadata", {}) or {} effective_backend = metadata.get("effective_backend", getattr(self.model, "lora_backend", None)) - if effective_backend != "fallback": + if effective_backend not in {"fallback", "peft"}: return 0 ema = getattr(getattr(self.trainer, "ema", None), "ema", None) if ema is None: return 0 - from ultralytics.utils.lora.fallback import FewShotLoRAConv, ManualLoRAConv online_modules = dict(self.model.named_modules()) ema_modules = dict(unwrap_model(ema).named_modules()) + if effective_backend == "peft": + return self._sync_peft_ema_scaling(online_modules, ema_modules) + + from ultralytics.utils.lora.fallback import FewShotLoRAConv, ManualLoRAConv + synced = 0 for name, online in online_modules.items(): if not isinstance(online, (ManualLoRAConv, FewShotLoRAConv)): @@ -251,6 +255,26 @@ def sync_ema_treatment(self) -> int: synced += 1 return synced + @staticmethod + def _sync_peft_ema_scaling(online_modules: dict, ema_modules: dict) -> int: + """Copy PEFT scaling dictionaries, which are intentionally absent from state_dict.""" + synced = 0 + for name, online in online_modules.items(): + if getattr(online, "lora_A", None) is None: + continue + averaged = ema_modules.get(name) + if averaged is None or getattr(averaged, "lora_A", None) is None: + raise ValueError(f"EMA PEFT adapter layout differs at '{name}'.") + online_scaling = getattr(online, "scaling", None) + ema_scaling = getattr(averaged, "scaling", None) + if not isinstance(online_scaling, dict): + continue + if not isinstance(ema_scaling, dict) or set(ema_scaling) != set(online_scaling): + raise ValueError(f"EMA PEFT scaling layout differs at '{name}'.") + ema_scaling.update(online_scaling) + synced += 1 + return synced + def _set_alpha_for_epoch(self, epoch: int) -> None: """Set the effective alpha schedule, including resume after the warmup endpoint.""" if self.strategy is None: diff --git a/ultralytics/nn/modules/moe/modules.py b/ultralytics/nn/modules/moe/modules.py index 1e5a3d08..8ae21297 100644 --- a/ultralytics/nn/modules/moe/modules.py +++ b/ultralytics/nn/modules/moe/modules.py @@ -404,7 +404,7 @@ class ES_MOE(nn.Module): def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8, top_k=2, use_sparse_inference=True, dynamic_threshold=0.4, - max_kernel_size=15): + max_kernel_size=15, expert_kernel_sizes=None): """ Args: in_channels: Input channels @@ -415,6 +415,10 @@ def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8, use_sparse_inference: Enable sparse Top-K expert computation during inference dynamic_threshold: Optional threshold for pruning low-confidence experts during inference max_kernel_size: Largest odd depthwise kernel assigned to an expert + expert_kernel_sizes: Optional explicit per-expert depthwise kernel sizes + (length must equal ``num_experts``). When ``None`` the kernels are + derived from defaults; pruned checkpoints set this so retraining + rebuilds the exact kept-expert kernels and reloads their weights. """ super(ES_MOE, self).__init__() @@ -440,6 +444,7 @@ def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8, self.in_channels = in_channels self.out_channels = out_channels self.num_experts = num_experts + self.reduction = reduction self.top_k = min(top_k, num_experts) if top_k is not None else num_experts self.use_top_k = (top_k is not None) self.use_sparse_inference = use_sparse_inference @@ -449,12 +454,27 @@ def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8, # Dynamic routing (Top-K supported) self.routing = DynamicRoutingLayer(in_channels, num_experts, reduction, top_k) - # Expert group (original design) - default_kernel_sizes = [3, 5, 7] - if num_experts <= len(default_kernel_sizes): - ks = [min(k, max_kernel_size) for k in default_kernel_sizes[:num_experts]] + # Expert group (original design). ``expert_kernel_sizes`` lets a pruned + # checkpoint reconstruct its kept experts' heterogeneous kernels so that + # ``YOLO(pruned.pt).train()`` reloads expert weights instead of dropping + # them on a kernel-shape mismatch (prune -> LoRA/full fine-tune recovery). + if expert_kernel_sizes is not None: + if len(expert_kernel_sizes) != num_experts: + raise ValueError( + f"expert_kernel_sizes must have {num_experts} entries, got {len(expert_kernel_sizes)}" + ) + ks = [] + for k in expert_kernel_sizes: + k = int(k) + if k % 2 == 0: + k -= 1 + ks.append(min(k, max_kernel_size)) else: - ks = [min(3 + 2 * i, max_kernel_size) for i in range(num_experts)] + default_kernel_sizes = [3, 5, 7] + if num_experts <= len(default_kernel_sizes): + ks = [min(k, max_kernel_size) for k in default_kernel_sizes[:num_experts]] + else: + ks = [min(3 + 2 * i, max_kernel_size) for i in range(num_experts)] self.experts = nn.ModuleList( [EfficientExpertGroup(in_channels, out_channels, kernel_size=k) for k in ks] ) diff --git a/ultralytics/nn/modules/moe/pruning.py b/ultralytics/nn/modules/moe/pruning.py index 913c4e79..963f9ab2 100644 --- a/ultralytics/nn/modules/moe/pruning.py +++ b/ultralytics/nn/modules/moe/pruning.py @@ -442,8 +442,30 @@ def _sync_yaml_num_experts(self, pruned_model: nn.Module) -> None: continue args = list(seq[j][3]) if isinstance(seq[j][3], list) else [seq[j][3]] out_ch = args[0] if args else getattr(mod, "out_channels", None) - seq[j][3] = [out_ch, int(mod.num_experts)] - LOGGER.info(" Synced post-prune expert counts into model.yaml") + # Recover the kept experts' actual depthwise kernel sizes; pruning keeps + # experts with heterogeneous kernels whose order/values are not the + # default [3, 5, 7...] a bare rebuild would assign. Writing the full + # positional arg list (incl. expert_kernel_sizes) lets YOLO(pruned.pt) + # .train() rebuild the exact experts so their weights survive + # intersect_dicts instead of being dropped on a kernel-shape mismatch. + kernels = [] + for expert in mod.experts: + conv = getattr(expert, "conv", None) + depthwise = getattr(conv, "depthwise", None) if conv is not None else None + ks = getattr(depthwise, "kernel_size", None) + kernels.append(int(ks[0]) if isinstance(ks, (tuple, list)) else int(ks) if ks else 3) + top_k = int(mod.top_k) if getattr(mod, "use_top_k", True) else None + seq[j][3] = [ + out_ch, + int(mod.num_experts), + int(getattr(mod, "reduction", 8)), + top_k, + bool(getattr(mod, "use_sparse_inference", True)), + float(getattr(mod, "dynamic_threshold", 0.4)), + int(getattr(mod, "max_kernel_size", 15)), + kernels, + ] + LOGGER.info(" Synced post-prune expert counts + kernel sizes into model.yaml") def _save_model(self, pruned_model: nn.Module, output_path: str) -> None: """