507 lines
18 KiB
Python
507 lines
18 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
YOLO11 关键点检测训练脚本（靶纸四角）。
|
||
|
||
设备优先级（--device auto）：Intel XPU > NVIDIA CUDA > CPU。
|
||
默认 imgsz=960；批大小默认 4（大图显存紧张时可再降）。
|
||
|
||
关于「业务像素误差」：
|
||
Ultralytics 没有在 yaml 里设定「像素阈值」的选项；反向传播仍由 pose/kobj/box 等内部 loss 驱动。
|
||
- 监控：--pixel-metrics-every N（每 N 个 epoch 打印 mean/p95，并合并进 runs/.../results.csv，见 pose_pixel_metrics.py）。
|
||
- 选 best.pt / early stopping：加 --best-by-pixel，用验证集 mean 像素误差（与 pose_pixel_metrics
|
||
同一口径）代替 mAP 合成 fitness（fitness = -mean_px，越小越好）。
|
||
多卡 DDP（world_size>1）时会自动退回默认 mAP fitness。
|
||
|
||
XPU：Ultralytics BaseTrainer._get_memory / _clear_memory 把非 MPS、非 CPU 一律当 CUDA，
|
||
会在验证前调用 torch.cuda 而报错；本脚本在选用 XPU 时自动打补丁（见 _patch_ultralytics_trainer_for_xpu）。
|
||
|
||
务必使用 pose 任务：YOLO(...) 与 model.train(...) 均指定 task='pose'。若误用默认 detect，
|
||
会把 17 列 Pose 标注当成检测/分割解析，校验时出现「coordinates > 1」或 [2.] 等假象。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import csv
|
||
import gc
|
||
import glob
|
||
import os
|
||
import tempfile
|
||
from copy import deepcopy
|
||
from pathlib import Path
|
||
|
||
import torch
|
||
from ultralytics import YOLO
|
||
|
||
from pose_pixel_metrics import eval_val_pixel_error
|
||
import warnings
|
||
warnings.filterwarnings('ignore',
|
||
message=".*scatter_add_kernel does not have a deterministic implementation.*")
|
||
|
||
|
||
def _clear_ultralytics_label_caches(data_yaml_path: str) -> int:
|
||
"""删除 data.yaml 的 path 下 labels/*.cache。
|
||
|
||
Ultralytics 的校验缓存 hash 仅依赖「标签/图片路径字符串 + 各文件 size 之和」，不含文件内容；
|
||
修正 *.txt 后若总和巧合不变，可能继续加载旧 cache 并重播旧的 corrupt 日志，训练前应删掉。"""
|
||
from ultralytics.utils import YAML
|
||
|
||
try:
|
||
cfg = YAML.load(data_yaml_path)
|
||
except Exception:
|
||
return 0
|
||
root = cfg.get("path")
|
||
if not root:
|
||
return 0
|
||
root = os.path.abspath(os.path.expanduser(str(root)))
|
||
pattern = os.path.join(root, "labels", "*.cache")
|
||
n = 0
|
||
for p in glob.glob(pattern):
|
||
try:
|
||
os.unlink(p)
|
||
n += 1
|
||
except OSError:
|
||
pass
|
||
return n
|
||
|
||
|
||
def _pick_device(explicit: str | None):
|
||
"""返回 ultralytics train/predict 可用的 device。"""
|
||
if explicit and explicit != "auto":
|
||
e = explicit.lower()
|
||
if e == "xpu":
|
||
if getattr(torch, "xpu", None) is None or not torch.xpu.is_available():
|
||
raise RuntimeError("指定了 --device xpu 但当前环境不可用")
|
||
return torch.device("xpu")
|
||
if e in ("0", "cuda", "gpu"):
|
||
if not torch.cuda.is_available():
|
||
raise RuntimeError("指定了 CUDA 但不可用")
|
||
return 0
|
||
if e == "cpu":
|
||
return "cpu"
|
||
return explicit
|
||
if getattr(torch, "xpu", None) is not None and torch.xpu.is_available():
|
||
return torch.device("xpu")
|
||
if torch.cuda.is_available():
|
||
return 0
|
||
return "cpu"
|
||
|
||
|
||
def _default_amp(device) -> bool:
|
||
if isinstance(device, torch.device) and device.type == "xpu":
|
||
return False
|
||
if device == "cpu":
|
||
return False
|
||
return True
|
||
|
||
|
||
def _patch_ultralytics_for_xpu():
|
||
"""为 Ultralytics 打补丁，使其能在 XPU 环境下正常训练和验证。"""
|
||
import ultralytics.engine.trainer as ut_trainer
|
||
import ultralytics.engine.validator as ut_validator
|
||
from ultralytics.utils.torch_utils import select_device as _original_select_device
|
||
|
||
# 1. 覆盖 select_device：Trainer 初始化传入 torch.device("xpu") 会走原版早返回；
|
||
# 初始化后 args.device 会变成字符串 "xpu"，中期 val 用 trainer.device，不调用 select_device；
|
||
# 训练结束 final_eval 里 Validator 会 select_device("xpu")，且 validator 在 import 时已绑定原函数，
|
||
# 只改 torch_utils 无效，必须同时修补 trainer/validator 模块内的引用。
|
||
def _patched_select_device(device="", *args, **kwargs):
|
||
# Ultralytics 8.4.x: select_device(device="", newline=False, verbose=True)
|
||
# Older forks sometimes passed extra positional args; forward everything.
|
||
if isinstance(device, str):
|
||
d = device.strip().lower()
|
||
if d == "xpu" or d.startswith("xpu:"):
|
||
return torch.device(device.strip())
|
||
return _original_select_device(device, *args, **kwargs)
|
||
|
||
import ultralytics.utils.torch_utils
|
||
|
||
ultralytics.utils.torch_utils.select_device = _patched_select_device
|
||
ut_trainer.select_device = _patched_select_device
|
||
ut_validator.select_device = _patched_select_device
|
||
|
||
# 2. 修补 Trainer 的内存函数
|
||
BT = ut_trainer.BaseTrainer
|
||
if not getattr(BT, "_archery_xpu_memory_patched", False):
|
||
_orig_get_memory = BT._get_memory
|
||
_orig_clear_memory = BT._clear_memory
|
||
|
||
def _get_memory(self, fraction=False):
|
||
if self.device.type != "xpu":
|
||
return _orig_get_memory(self, fraction)
|
||
# ... (原有的 XPU 内存获取逻辑保持不变) ...
|
||
memory, total = 0, 0
|
||
try:
|
||
idx = self.device.index
|
||
if idx is None:
|
||
idx = torch.xpu.current_device()
|
||
memory = int(torch.xpu.memory_allocated(idx))
|
||
if fraction:
|
||
total = int(torch.xpu.get_device_properties(idx).total_memory)
|
||
except Exception:
|
||
pass
|
||
return (memory / total) if fraction and total > 0 else (memory / 2**30)
|
||
|
||
def _clear_memory(self, threshold=None):
|
||
if self.device.type != "xpu":
|
||
return _orig_clear_memory(self, threshold)
|
||
if threshold is not None:
|
||
assert 0 <= threshold <= 1, "Threshold must be between 0 and 1."
|
||
if self._get_memory(fraction=True) <= threshold:
|
||
return
|
||
gc.collect()
|
||
if hasattr(torch.xpu, "empty_cache"):
|
||
torch.xpu.empty_cache()
|
||
|
||
BT._get_memory = _get_memory
|
||
BT._clear_memory = _clear_memory
|
||
BT._archery_xpu_memory_patched = True
|
||
|
||
# 3. 修补 Validator 的内存函数 (关键是添加这部分)
|
||
BV = ut_validator.BaseValidator
|
||
if not getattr(BV, "_archery_xpu_memory_patched", False):
|
||
# 为 Validator 添加同样的内存处理方法
|
||
BV._get_memory = _get_memory
|
||
BV._clear_memory = _clear_memory
|
||
BV._archery_xpu_memory_patched = True
|
||
|
||
|
||
def _install_best_by_pixel_validate(data_yaml: str, imgsz: int, conf: float) -> None:
|
||
"""用验证集关键点像素 mean 替代 mAP fitness，驱动 best.pt 与 patience early stopping。"""
|
||
import ultralytics.engine.trainer as ut
|
||
from ultralytics.utils import RANK
|
||
|
||
BT = ut.BaseTrainer
|
||
if getattr(BT, "_archery_best_by_pixel_installed", False):
|
||
return
|
||
|
||
_orig_validate = BT.validate
|
||
|
||
def validate(self):
|
||
import torch.distributed as dist
|
||
|
||
if self.ema and self.world_size > 1:
|
||
for buffer in self.ema.ema.buffers():
|
||
dist.broadcast(buffer, src=0)
|
||
metrics = self.validator(self)
|
||
if metrics is None:
|
||
return None, None
|
||
orig_fitness = metrics.pop("fitness", -self.loss.detach().cpu().numpy())
|
||
|
||
use_pixel = self.world_size <= 1 and RANK in {-1, 0}
|
||
mean_px: float | None = None
|
||
if use_pixel:
|
||
tmp_path: str | None = None
|
||
try:
|
||
fd, tmp_path = tempfile.mkstemp(suffix=".pt", prefix="archery_pxfit_")
|
||
os.close(fd)
|
||
from ultralytics.utils.torch_utils import unwrap_model
|
||
|
||
core = unwrap_model(self.ema.ema if self.ema else self.model)
|
||
torch.save({"ema": deepcopy(core).half(), "train_args": vars(self.args)}, tmp_path)
|
||
probe = YOLO(tmp_path)
|
||
stats = eval_val_pixel_error(
|
||
probe,
|
||
data_yaml,
|
||
device=self.device,
|
||
imgsz=imgsz,
|
||
conf=conf,
|
||
)
|
||
mean_px = stats.get("mean_px")
|
||
if mean_px is None:
|
||
raise RuntimeError("无有效 mean_px（检查 val 标签与检测是否为空）")
|
||
except Exception as exc:
|
||
print(f"\n⚠️ [best-by-pixel] 像素探针失败，本 epoch 仍用 mAP fitness: {exc}\n")
|
||
mean_px = None
|
||
finally:
|
||
if tmp_path:
|
||
try:
|
||
os.unlink(tmp_path)
|
||
except OSError:
|
||
pass
|
||
|
||
if mean_px is not None:
|
||
fitness = -float(mean_px)
|
||
metrics["metrics/mean_px(val)"] = float(mean_px)
|
||
else:
|
||
fitness = float(orig_fitness)
|
||
|
||
if not self.best_fitness or self.best_fitness < fitness:
|
||
self.best_fitness = fitness
|
||
return metrics, fitness
|
||
|
||
BT.validate = validate
|
||
BT._archery_best_by_pixel_installed = True
|
||
|
||
|
||
def _fmt_csv_metric(v: float | int | None) -> str:
|
||
if v is None:
|
||
return ""
|
||
if isinstance(v, float):
|
||
return f"{v:.6g}"
|
||
return str(v)
|
||
|
||
|
||
# 写入 results.csv 的列名（与 --best-by-pixel 的 metrics/mean_px(val) 区分，避免被 last.pt 回调覆盖 EMA 行）
|
||
_PIXEL_METRIC_COLUMNS: tuple[tuple[str, str], ...] = (
|
||
("pixel_error/mean_px", "mean_px"),
|
||
("pixel_error/median_px", "median_px"),
|
||
("pixel_error/p95_px", "p95_px"),
|
||
("pixel_error/max_px", "max_px"),
|
||
("pixel_error/n_points", "n_points"),
|
||
("pixel_error/n_images", "n_images"),
|
||
("pixel_error/skip_no_det", "skip_no_det"),
|
||
("pixel_error/skip_no_gt", "skip_no_gt"),
|
||
("pixel_error/skip_kpt_mismatch", "skip_kpt_mismatch"),
|
||
)
|
||
|
||
|
||
def _merge_pixel_metrics_into_results_csv(save_dir: str | Path, epoch_1based: int, stats: dict) -> None:
|
||
"""在 Ultralytics 写完本 epoch 行之后，把像素指标列合并进 results.csv（扩展表头、补空列）。"""
|
||
csv_path = Path(save_dir) / "results.csv"
|
||
if not csv_path.is_file():
|
||
return
|
||
try:
|
||
with open(csv_path, newline="", encoding="utf-8") as f:
|
||
rows = list(csv.reader(f))
|
||
except OSError:
|
||
return
|
||
if len(rows) < 2:
|
||
return
|
||
header = list(rows[0])
|
||
for col_name, _ in _PIXEL_METRIC_COLUMNS:
|
||
if col_name not in header:
|
||
header.append(col_name)
|
||
for ri in range(1, len(rows)):
|
||
rows[ri].append("")
|
||
col_ix = {name: i for i, name in enumerate(header)}
|
||
rows[0] = header
|
||
target_ri: int | None = None
|
||
for ri in range(1, len(rows)):
|
||
row = rows[ri]
|
||
while len(row) < len(header):
|
||
row.append("")
|
||
try:
|
||
if int(float(row[0].strip())) == int(epoch_1based):
|
||
target_ri = ri
|
||
except (ValueError, IndexError):
|
||
continue
|
||
if target_ri is None:
|
||
return
|
||
row = rows[target_ri]
|
||
while len(row) < len(header):
|
||
row.append("")
|
||
for col_name, sk in _PIXEL_METRIC_COLUMNS:
|
||
row[col_ix[col_name]] = _fmt_csv_metric(stats.get(sk))
|
||
try:
|
||
with open(csv_path, "w", newline="", encoding="utf-8") as f:
|
||
w = csv.writer(f)
|
||
w.writerows(rows)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _make_pixel_metrics_callback(data_yaml: str, every: int, imgsz: int, conf: float = 0.25):
|
||
def on_fit_epoch_end(trainer):
|
||
from ultralytics.utils import RANK
|
||
|
||
if RANK not in {-1, 0}:
|
||
return
|
||
if every <= 0:
|
||
return
|
||
ep = int(getattr(trainer, "epoch", -1))
|
||
if (ep + 1) % every != 0:
|
||
return
|
||
w = Path(trainer.save_dir) / "weights" / "last.pt"
|
||
if not w.is_file():
|
||
return
|
||
m = YOLO(str(w))
|
||
stats = eval_val_pixel_error(
|
||
m,
|
||
data_yaml,
|
||
device=trainer.device,
|
||
imgsz=imgsz,
|
||
conf=conf,
|
||
)
|
||
mean_px = stats.get("mean_px")
|
||
p95_px = stats.get("p95_px")
|
||
mean_s = f"{mean_px:.3f}" if mean_px is not None else "n/a"
|
||
p95_s = f"{p95_px:.3f}" if p95_px is not None else "n/a"
|
||
print(
|
||
f"\n[pixel-metrics] epoch {ep + 1}: mean_px={mean_s} p95_px={p95_s} "
|
||
f"n_points={stats.get('n_points', 0)} "
|
||
f"skip(det/gt/k)={stats['skip_no_det']}/{stats['skip_no_gt']}/{stats['skip_kpt_mismatch']}\n"
|
||
)
|
||
_merge_pixel_metrics_into_results_csv(trainer.save_dir, ep + 1, stats)
|
||
|
||
return on_fit_epoch_end
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(description="YOLO Pose 训练（XPU/CUDA/CPU）")
|
||
ap.add_argument("--data", default="datasets/dataset_pose.yaml", help="data.yaml")
|
||
ap.add_argument("--model", default="yolo11x-pose.pt", help="预训练权重")
|
||
ap.add_argument("--epochs", type=int, default=100)
|
||
ap.add_argument("--imgsz", type=int, default=960, help="训练输入边长（默认 960）")
|
||
ap.add_argument("--batch", type=int, default=4, help="批大小；OOM 时减小")
|
||
ap.add_argument(
|
||
"--device",
|
||
default="auto",
|
||
help="auto | xpu | 0 | cuda | cpu（auto：XPU 优先）",
|
||
)
|
||
ap.add_argument(
|
||
"--no-amp",
|
||
action="store_true",
|
||
help="关闭混合精度（默认：CUDA 开启，XPU/CPU 关闭）",
|
||
)
|
||
ap.add_argument("--project", default="runs/pose")
|
||
ap.add_argument("--name", default="target_pose_train")
|
||
ap.add_argument("--workers", type=int, default=4)
|
||
ap.add_argument(
|
||
"--pixel-metrics-every",
|
||
type=int,
|
||
default=0,
|
||
help="每 N 个 epoch 在 val 上打印像素误差并写入 results.csv 对应 epoch 行（0=关闭）；需 labels 与 data.yaml 布局一致",
|
||
)
|
||
ap.add_argument(
|
||
"--pixel-metrics-conf",
|
||
type=float,
|
||
default=0.25,
|
||
help="--pixel-metrics-every 时 predict 置信度阈值（默认 0.25）",
|
||
)
|
||
ap.add_argument(
|
||
"--best-by-pixel",
|
||
action="store_true",
|
||
help="best.pt 与 early stopping 按验证集 mean 像素误差（同 pose_pixel_metrics），fitness=-mean_px；单卡有效，DDP 自动退回 mAP",
|
||
)
|
||
ap.add_argument(
|
||
"--pixel-fitness-conf",
|
||
type=float,
|
||
default=0.25,
|
||
help="--best-by-pixel 时 predict 置信度阈值（默认与 pixel-metrics 一致）",
|
||
)
|
||
ap.add_argument(
|
||
"--export-onnx",
|
||
action="store_true",
|
||
help="训练结束后导出 ONNX（需再设 --onnx-imgsz）",
|
||
)
|
||
ap.add_argument(
|
||
"--onnx-imgsz",
|
||
type=int,
|
||
nargs=2,
|
||
metavar=("H", "W"),
|
||
default=[224, 320],
|
||
help="导出 ONNX 的 [高, 宽]，默认 224 320（Maix 常用）",
|
||
)
|
||
ap.add_argument(
|
||
"--clear-label-cache",
|
||
action="store_true",
|
||
help="启动训练前删除 data.yaml 中 path 下的 labels/*.cache（修正标注后仍报 corrupt 时用）",
|
||
)
|
||
args = ap.parse_args()
|
||
|
||
device = _pick_device(None if args.device == "auto" else args.device)
|
||
use_amp = False if args.no_amp else _default_amp(device)
|
||
|
||
if isinstance(device, torch.device) and device.type == "xpu":
|
||
print(f"✅ 使用 Intel XPU: {device}")
|
||
elif device == 0 or device == "0":
|
||
print(f"✅ 使用 CUDA: {torch.cuda.get_device_name(0)}")
|
||
else:
|
||
print("⚠️ 使用 CPU，训练会较慢")
|
||
|
||
if isinstance(device, torch.device) and device.type == "xpu":
|
||
_patch_ultralytics_for_xpu()
|
||
|
||
data_yaml = args.data
|
||
if not os.path.isabs(data_yaml):
|
||
data_yaml = os.path.join(os.path.dirname(os.path.abspath(__file__)), data_yaml)
|
||
if not os.path.exists(data_yaml):
|
||
print(f"❌ 数据集配置不存在: {data_yaml}")
|
||
return
|
||
|
||
if args.clear_label_cache:
|
||
n_rm = _clear_ultralytics_label_caches(data_yaml)
|
||
print(f"🗑️ 已删除标签目录缓存 {n_rm} 个（labels/*.cache），将强制重新扫描标注。")
|
||
|
||
print(f"📦 加载模型: {args.model}（固定 task=pose）")
|
||
model = YOLO(args.model, task="pose")
|
||
|
||
if args.best_by_pixel:
|
||
_install_best_by_pixel_validate(data_yaml, args.imgsz, args.pixel_fitness_conf)
|
||
print(
|
||
"📌 已启用 --best-by-pixel：best.pt / patience 按验证集 mean 像素误差（fitness=-mean_px）；"
|
||
"反向传播仍为 Ultralytics 默认 pose/box loss。"
|
||
)
|
||
|
||
if args.pixel_metrics_every > 0:
|
||
model.add_callback(
|
||
"on_fit_epoch_end",
|
||
_make_pixel_metrics_callback(
|
||
data_yaml, args.pixel_metrics_every, args.imgsz, conf=args.pixel_metrics_conf
|
||
),
|
||
)
|
||
|
||
model.train(
|
||
task="pose",
|
||
data=data_yaml,
|
||
epochs=args.epochs,
|
||
imgsz=args.imgsz,
|
||
batch=args.batch,
|
||
name=args.name,
|
||
project=args.project,
|
||
exist_ok=True,
|
||
save=True,
|
||
save_period=5,
|
||
device=device,
|
||
workers=args.workers,
|
||
lr0=0.0001,
|
||
lrf=0.01,
|
||
optimizer="AdamW",
|
||
momentum=0.937,
|
||
weight_decay=0.001,
|
||
warmup_epochs=0,
|
||
warmup_momentum=0.8,
|
||
warmup_bias_lr=0.1,
|
||
hsv_h=0.015,
|
||
hsv_s=0.7,
|
||
hsv_v=0.4,
|
||
degrees=5.0,
|
||
translate=0.0,
|
||
scale=0.2,
|
||
shear=0.0,
|
||
perspective=0.0000,
|
||
flipud=0.0,
|
||
fliplr=0.5,
|
||
mosaic=0.0,
|
||
mixup=0.0,
|
||
copy_paste=0.0,
|
||
box=6,
|
||
cls=0.5,
|
||
dfl=1.5,
|
||
pose=18.0,
|
||
kobj=0.5,
|
||
freeze=0,
|
||
seed=42,
|
||
verbose=True,
|
||
amp=use_amp,
|
||
patience=100,
|
||
cos_lr=True,
|
||
)
|
||
|
||
print("\n✅ 训练完成！")
|
||
print(f"📁 best: {args.project}/{args.name}/weights/best.pt")
|
||
print(f"📁 last: {args.project}/{args.name}/weights/last.pt")
|
||
print("📊 仅看像素误差可运行: python pose_pixel_metrics.py --model <best.pt> --data <yaml> --imgsz", args.imgsz)
|
||
|
||
if args.export_onnx:
|
||
h, w = args.onnx_imgsz
|
||
print(f"📦 导出 ONNX imgsz=[{h}, {w}] ...")
|
||
model.export(format="onnx", imgsz=[h, w], simplify=True, opset=17, dynamic=False)
|
||
print("✅ ONNX 完成")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|