diff --git a/README_reproduce.md b/README_reproduce.md new file mode 100644 index 00000000..a7098ccb --- /dev/null +++ b/README_reproduce.md @@ -0,0 +1,75 @@ + + +# YOLO-Master 复现：VisDrone & SKU-110K + +本目录包含在 VisDrone 和 SKU-110K 数据集上复现 YOLO-Master v0.1-N 和 EsMoE-N 的训练脚本与结果。 + +## 数据集下载 +- *VisDrone2019-DET: [官网下载](http://aiskyeye.com/)，解压后按 `ultralytics/cfg/datasets/VisDrone.yaml` 组织目录。 +- *SKU-110K: [GitHub 项目页](https://github.com/eg4000/SKU110K_CVPR19) 下载，解压后按 `ultralytics/cfg/datasets/SKU-110K.yaml` 组织目录。 + +## 训练命令 +### VisDrone +```bash +# 训练 v0.1-N 基线 +python scripts/reproduce/reproduce_visdrone.py --model v01 + +# 训练 EsMoE-N +python scripts/reproduce/reproduce_visdrone.py --model moe +``` + +### SKU-110K +```bash +# 训练 v0.1-N 基线 +python scripts/reproduce/reproduce_sku110k.py --model v01 + +# 训练 EsMoE-N +python scripts/reproduce/reproduce_sku110k.py --model moe +``` + +## 复现结果对比 + +| 数据集 | 模型 | 输入分辨率 | 训练轮数 | mAP50 | mAP50-95 | 参数量 | +|--------|------|------------|----------|-------|----------|--------| +| VisDrone | YOLO-Master v0.1-N | 800 | 120 | 0.360 | 0.213 | 7.5M | +| VisDrone | YOLO-Master EsMoE-N | 800 | 120 | 0.360 | 0.212 | 3.4M | +| SKU-110K | YOLO-Master v0.1-N | 640 | 120 | 0.885 | 0.564 | 7.5M | +| SKU-110K | YOLO-Master EsMoE-N | 640 | 120 | 0.886 | 0.563 | 3.4M | + +## 训练日志查看 +### 离线日志下载（百度网盘） +由于 WandB 项目设置为私有，无法直接公开访问，我们将所有离线运行包上传至百度网盘： +- **下载链接**: [https://pan.baidu.com/s/1CG9I9rxe-Z2Pmhla06fnbg?pwd=tiw2](https://pan.baidu.com/s/1CG9I9rxe-Z2Pmhla06fnbg?pwd=tiw2) + **提取码**: `tiw2` +- 文件夹内包含以下离线日志： + - `wandb_visdrone_v01n_800.zip` + - `wandb_visdrone_moe_800.zip` + - `wandb_sku110k_v01n_640.zip` + - `wandb_sku110k_moe_640.zip` + +### 使用方法 +1. 下载所需 zip 文件并解压。 +2. 在命令行执行（以 VisDrone v0.1-N 为例，文件夹名请以实际解压结果为准）： + ```bash + wandb sync ./run-xxx.wandb + ``` +3. 浏览器会自动打开 WandB 页面，显示完整的损失曲线、mAP 曲线等所有指标。 + +## 已知问题 +- WandB 为私有团队空间，无法提供可公开访问的在线链接，请使用离线包。 +- Windows 训练时必须设置 `workers=0` 以避免多进程死锁。 +- 若显存不足（<8GB），请减小脚本中的 `BATCH` 变量（例如 VisDrone 可改为 4，SKU-110K 改为 8）。 +- VisDrone 小物体密集，`imgsz=800` 为推荐最低值，不宜进一步降低。 + +## 本地复现环境 +- GPU: NVIDIA RTX 5060 Laptop (8GB) +- CUDA Driver: 13.1 +- PyTorch: 2.11.0 (CUDA 12.8) +- 依赖安装命令： + ```bash + pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu128 + pip install ultralytics wandb + ``` +``` + + diff --git a/scripts/reproduce/reproduce_sku110k.py b/scripts/reproduce/reproduce_sku110k.py index 69b6640d..5e8165db 100644 --- a/scripts/reproduce/reproduce_sku110k.py +++ b/scripts/reproduce/reproduce_sku110k.py @@ -1,33 +1,167 @@ -#!/usr/bin/env python3 -"""Reproduce YOLO-Master-v0.1-N and YOLO-Master-EsMoE-N baselines on SKU-110K. - -SKU-110K (retail, dense products, single class), built-in config SKU-110K.yaml. -By default the models are reproduced as-is (EsMoE-N keeps its sparse eval, which -collapses mAP). Add --no-sparse-eval to opt into the corrected dense evaluation -for EsMoE-N (train==eval); v0.1-N is unaffected. - -Examples: - python scripts/reproduce/reproduce_sku110k.py --check-build - python scripts/reproduce/reproduce_sku110k.py --epochs 300 --batch 64 # as-is - python scripts/reproduce/reproduce_sku110k.py --model EsMoE-N --no-sparse-eval # corrected - python scripts/reproduce/reproduce_sku110k.py --model v0.1-N --no-wandb - python scripts/reproduce/reproduce_sku110k.py --wandb-project my-proj --wandb-mode offline -""" -from __future__ import annotations - -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parent)) - -from _reproduce_common import DatasetSpec, run_dataset # noqa: E402 - -DATASET = DatasetSpec( - name="SKU-110K", - data="SKU-110K.yaml", - project="runs/reproduce/sku110k", -) - - -if __name__ == "__main__": - raise SystemExit(run_dataset(DATASET)) +#!/usr/bin/env python +# -*- coding: utf-8 -*- +""" +复现 SKU-110K：YOLO-Master v0.1-N vs EsMoE-N +单模型独立训练，强制 WandB 在线，数据完整保留 + +用法: + python scripts/reproduce/reproduce_sku110k.py --model v01 # 训练 v0.1-N + python scripts/reproduce/reproduce_sku110k.py --model moe # 训练 EsMoE-N +""" + +import os +import sys +import argparse +import traceback + +ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, ROOT) + +import wandb +from ultralytics import YOLO + +# ================== 环境 & WandB 强制在线 ================== +os.environ["WANDB_MODE"] = "online" +os.environ["WANDB_SERVICE_WAIT"] = "300" +os.environ["WANDB_INIT_TIMEOUT"] = "300" +os.environ["WANDB__SERVICE_WAIT"] = "300" +os.environ["WANDB_START_METHOD"] = "thread" +os.environ["WANDB_DISABLE_SERVICE"] = "True" +os.environ["WANDB_SILENT"] = "false" +os.environ["WANDB_LOG_MODEL"] = "false" + +# ================== 固定超参 ================== +DATA_YAML = "ultralytics/cfg/datasets/SKU-110K.yaml" +IMG_SIZE = 640 +EPOCHS = 120 +BATCH = 12 # 8G 显存安全值，若 OOM 改为 8 +DEVICE = "0" +PROJECT = "SKU110K_Reproduce" +NAME_BASE = "SKU110K" + +NOTE_V01 = ( + "YOLO-Master v0.1-N on SKU-110K; " + "imgsz=640, batch=12, amp, cache=disk" +) +NOTE_MOE = ( + "YOLO-Master EsMoE-N on SKU-110K; " + "imgsz=640, batch=12, amp, cache=disk" +) + +def ensure_wandb_login(): + """确保 WandB 可用，无 API key 则用匿名模式""" + try: + if wandb.api.api_key: + print("✅ WandB API key 已配置") + return True + except Exception: + pass + key = os.environ.get("WANDB_API_KEY") + if key: + wandb.login(key=key) + return True + print("⚠️ 未检测到 API key，使用匿名模式（数据仍保留，链接公开）") + wandb.init(anonymous="allow", project=PROJECT, mode="online") + wandb.finish() + return False + +def init_run(run_name, notes): + """初始化 WandB run，带超时保护""" + for attempt in range(3): + try: + run = wandb.init( + project=PROJECT, + name=run_name, + notes=notes, + resume="allow", + config={"imgsz": IMG_SIZE, "batch": BATCH, "epochs": EPOCHS}, + settings=wandb.Settings(_service_wait=300, init_timeout=120) + ) + print(f"✅ WandB Run 初始化成功, ID: {run.id}") + return run + except Exception as e: + print(f"❌ 初始化失败 (attempt {attempt+1}/3): {e}") + if attempt < 2: + import time + time.sleep(10) + else: + raise RuntimeError("无法连接 WandB，请检查网络或 API key") from e + +def train(model_cfg, note, run_name): + """单次训练流程（针对大分辨率数据集优化）""" + print(f"\n{'='*60}") + print(f"🚀 正在训练: {run_name}") + print(f"📁 数据: {DATA_YAML}") + print(f"⚙️ 模型: {model_cfg}") + print(f"🍃 参数: imgsz={IMG_SIZE}, batch={BATCH}, epochs={EPOCHS}") + print(f"{'='*60}") + + ensure_wandb_login() + run = init_run(run_name, note) + + try: + model = YOLO(model_cfg) + model.train( + data=DATA_YAML, + epochs=EPOCHS, + imgsz=IMG_SIZE, + batch=BATCH, + device=DEVICE, + amp=True, + # ✅ 修改点1：使用磁盘缓存，避免内存爆炸 + cache="disk", + # ✅ 修改点2：Windows 下 workers=0 最稳定，避免 I/O 死锁 + workers=0, + # ✅ 修改点3：矩形训练，减少 padding 浪费，加速且省显存 + rect=True, + # ✅ 修改点4：自动调整 batch size + nbs=64, + # ✅ 修改点5：最后10个epoch关闭mosaic，提升最终精度 + close_mosaic=10, + project=PROJECT, + name=run_name, + exist_ok=True, + save=True, + save_period=10, + plots=True, + # 这两个保存选项会拖慢速度，复现实验可暂时关闭 + # save_json=True, + # save_hybrid=True, + ) + print(f"✅ {run_name} 训练完成") + + # 最终验证 + val_results = model.val(data=DATA_YAML, split='val') + print(f"📊 验证结果: mAP50={val_results.box.map50:.4f}, mAP50-95={val_results.box.map:.4f}") + except Exception as e: + print(f"❌ 训练过程中出错: {e}") + traceback.print_exc() + finally: + print("⏳ 同步 WandB 数据...") + wandb.finish() + print("✅ WandB Run 已结束，数据已完整上传。") + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="SKU-110K 复现训练") + parser.add_argument( + "--model", type=str, required=True, + choices=["v01", "moe"], + help="选择要训练的模型: v01 (YOLO-Master v0.1-N) 或 moe (EsMoE-N)" + ) + args = parser.parse_args() + + if args.model == "v01": + train( + model_cfg="ultralytics/cfg/models/master/v0_1/det/yolo-master-n.yaml", + note=NOTE_V01, + run_name=f"{NAME_BASE}_v01_640_ep{EPOCHS}" + ) + else: # moe + train( + model_cfg="ultralytics/cfg/models/master/exp/yolo-master-v0_10.yaml", + note=NOTE_MOE, + run_name=f"{NAME_BASE}_MoE_640_ep{EPOCHS}" + ) + + print("\n🎉 训练完毕。前往 https://wandb.ai 查看实验。") + print("📘 将项目设置为 Public 即可获取公开对比链接。") \ No newline at end of file diff --git a/scripts/reproduce/reproduce_visdrone.py b/scripts/reproduce/reproduce_visdrone.py index d987b199..fdd1bbc2 100644 --- a/scripts/reproduce/reproduce_visdrone.py +++ b/scripts/reproduce/reproduce_visdrone.py @@ -1,33 +1,170 @@ -#!/usr/bin/env python3 -"""Reproduce YOLO-Master-v0.1-N and YOLO-Master-EsMoE-N baselines on VisDrone. - -VisDrone (aerial, dense small objects), built-in config VisDrone.yaml. -By default the models are reproduced as-is (EsMoE-N keeps its sparse eval, which -collapses mAP). Add --no-sparse-eval to opt into the corrected dense evaluation -for EsMoE-N (train==eval); v0.1-N is unaffected. - -Examples: - python scripts/reproduce/reproduce_visdrone.py --check-build - python scripts/reproduce/reproduce_visdrone.py --epochs 300 --batch 64 # as-is - python scripts/reproduce/reproduce_visdrone.py --model EsMoE-N --no-sparse-eval # corrected - python scripts/reproduce/reproduce_visdrone.py --model v0.1-N --no-wandb - python scripts/reproduce/reproduce_visdrone.py --wandb-project my-proj --wandb-mode offline -""" -from __future__ import annotations - -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parent)) - -from _reproduce_common import DatasetSpec, run_dataset # noqa: E402 - -DATASET = DatasetSpec( - name="VisDrone", - data="VisDrone.yaml", - project="runs/reproduce/visdrone", -) - - -if __name__ == "__main__": - raise SystemExit(run_dataset(DATASET)) +#!/usr/bin/env python +# -*- coding: utf-8 -*- +""" +复现 VisDrone2019：YOLO-Master v0.1-N vs EsMoE-N +单模型独立训练，强制 WandB 在线，数据完整保留 + +用法: + python scripts/reproduce/reproduce_visdrone.py --model v01 # 训练 v0.1-N + python scripts/reproduce/reproduce_visdrone.py --model moe # 训练 EsMoE-N +""" + +import os +import sys +import argparse +import traceback +from multiprocessing import freeze_support + +ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, ROOT) + +import wandb +from ultralytics import YOLO + +# ================== 环境 & WandB 强制在线 ================== +os.environ["WANDB_MODE"] = "online" +os.environ["WANDB_SERVICE_WAIT"] = "300" +os.environ["WANDB_INIT_TIMEOUT"] = "300" +os.environ["WANDB__SERVICE_WAIT"] = "300" +os.environ["WANDB_START_METHOD"] = "thread" +os.environ["WANDB_DISABLE_SERVICE"] = "True" +os.environ["WANDB_SILENT"] = "false" +os.environ["WANDB_LOG_MODEL"] = "false" + +# ================== 固定超参（针对 VisDrone 优化）================== +DATA_YAML = "ultralytics/cfg/datasets/VisDrone.yaml" +IMG_SIZE = 800 # VisDrone 推荐较高分辨率以保留小目标 +EPOCHS = 120 +BATCH = 6 # 8G 显存安全值，若 OOM 可改为 4 +DEVICE = "0" +PROJECT = "VisDrone_Reproduce" +NAME_BASE = "VisDrone" + +NOTE_V01 = ( + "YOLO-Master v0.1-N on VisDrone2019; " + "imgsz=800, batch=6, amp, cache=disk" +) +NOTE_MOE = ( + "YOLO-Master EsMoE-N on VisDrone2019; " + "imgsz=800, batch=6, amp, cache=disk" +) + +def ensure_wandb_login(): + """确保 WandB 可用，无 API key 则用匿名模式""" + try: + if wandb.api.api_key: + print("✅ WandB API key 已配置") + return True + except Exception: + pass + key = os.environ.get("WANDB_API_KEY") + if key: + wandb.login(key=key) + return True + print("⚠️ 未检测到 API key，使用匿名模式（数据仍保留，链接公开）") + wandb.init(anonymous="allow", project=PROJECT, mode="online") + wandb.finish() + return False + +def init_run(run_name, notes): + """初始化 WandB run，带超时保护""" + for attempt in range(3): + try: + run = wandb.init( + project=PROJECT, + name=run_name, + notes=notes, + resume="allow", + config={"imgsz": IMG_SIZE, "batch": BATCH, "epochs": EPOCHS}, + settings=wandb.Settings(_service_wait=300, init_timeout=120) + ) + print(f"✅ WandB Run 初始化成功, ID: {run.id}") + return run + except Exception as e: + print(f"❌ 初始化失败 (attempt {attempt+1}/3): {e}") + if attempt < 2: + import time + time.sleep(10) + else: + raise RuntimeError("无法连接 WandB，请检查网络或 API key") from e + +def train(model_cfg, note, run_name): + """单次训练流程（针对 VisDrone 密集小目标优化）""" + print(f"\n{'='*60}") + print(f"🚀 正在训练: {run_name}") + print(f"📁 数据: {DATA_YAML}") + print(f"⚙️ 模型: {model_cfg}") + print(f"🍃 参数: imgsz={IMG_SIZE}, batch={BATCH}, epochs={EPOCHS}") + print(f"{'='*60}") + + ensure_wandb_login() + run = init_run(run_name, note) + + try: + model = YOLO(model_cfg) + model.train( + data=DATA_YAML, + epochs=EPOCHS, + imgsz=IMG_SIZE, + batch=BATCH, + device=DEVICE, + amp=True, + # 使用磁盘缓存，避免内存爆炸 + cache="disk", + # Windows 下 workers=0 最稳定，避免 I/O 死锁 + workers=0, + # 数据增强（复现原始 VisDrone 脚本中的经验设置） + mosaic=1.0, + mixup=0.1, + copy_paste=0.1, + # 最后 15 个 epoch 关闭 mosaic，稳定最终精度 + close_mosaic=15, + # 早停，避免无意义训练 + patience=20, + # 不启用矩形训练（VisDrone 图像尺寸差异大，rect 无益） + rect=False, + project=PROJECT, + name=run_name, + exist_ok=True, + save=True, + save_period=10, + plots=True, + ) + print(f"✅ {run_name} 训练完成") + + # 最终验证 + val_results = model.val(data=DATA_YAML, split='val') + print(f"📊 验证结果: mAP50={val_results.box.map50:.4f}, mAP50-95={val_results.box.map:.4f}") + except Exception as e: + print(f"❌ 训练过程中出错: {e}") + traceback.print_exc() + finally: + print("⏳ 同步 WandB 数据...") + wandb.finish() + print("✅ WandB Run 已结束，数据已完整上传。") + +if __name__ == "__main__": + freeze_support() + parser = argparse.ArgumentParser(description="VisDrone 复现训练") + parser.add_argument( + "--model", type=str, required=True, + choices=["v01", "moe"], + help="选择要训练的模型: v01 (YOLO-Master v0.1-N) 或 moe (EsMoE-N)" + ) + args = parser.parse_args() + + if args.model == "v01": + train( + model_cfg="ultralytics/cfg/models/master/v0_1/det/yolo-master-n.yaml", + note=NOTE_V01, + run_name=f"{NAME_BASE}_v01_800_ep{EPOCHS}" + ) + else: # moe + train( + model_cfg="ultralytics/cfg/models/master/exp/yolo-master-v0_10.yaml", + note=NOTE_MOE, + run_name=f"{NAME_BASE}_MoE_800_ep{EPOCHS}" + ) + + print("\n🎉 训练完毕。前往 https://wandb.ai 查看实验。") + print("📘 将项目设置为 Public 即可获取公开对比链接。") \ No newline at end of file