377 lines
11 KiB
Python
377 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
自动评估 FUSIONLCD 多个 checkpoint 的脚本
|
||
|
||
支持：
|
||
1. 单卡串行测试：
|
||
--gpu 0
|
||
2. 多卡并行测试：
|
||
--gpu 0,1,2,3
|
||
|
||
python auto_eval_checkpoints.py \
|
||
--project_dir /home/adlab36/chenyouyuan/FUSIONLCD \
|
||
--config /home/adlab36/chenyouyuan/FUSIONLCD/config.yaml \
|
||
--train_script /home/adlab36/chenyouyuan/FUSIONLCD/train.py \
|
||
--models_dir /home/adlab36/chenyouyuan/FUSIONLCD/result/log/models \
|
||
--result_name auto_eval \
|
||
--gpu 2,3 \
|
||
--epochs_filter 119, 139
|
||
|
||
099, 119, 139, 159, 179, 199
|
||
说明：
|
||
- 多卡模式下，每个 checkpoint 会分配到一个 GPU
|
||
- 每个子进程使用独立的临时工作目录和独立 config.yaml，避免冲突
|
||
- 会实时输出子进程日志
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import csv
|
||
import os
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
import tempfile
|
||
from pathlib import Path
|
||
from typing import Dict, List, Tuple, Optional
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
|
||
import yaml
|
||
|
||
|
||
CKPT_RE = re.compile(r"checkpoint_(\d+)\.pth\.tar$")
|
||
RESULT_LINE_RE = re.compile(r"^\d{14}\s+(\d+)\s+(\d+)\s+(.*)$")
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--project_dir", type=str, required=True, help="项目根目录")
|
||
parser.add_argument("--config", type=str, required=True, help="config.yaml 路径")
|
||
parser.add_argument("--train_script", type=str, required=True, help="train.py 路径")
|
||
parser.add_argument("--models_dir", type=str, required=True, help="checkpoint 目录")
|
||
parser.add_argument("--result_name", type=str, default="auto_eval", help="train.py 的 result_name")
|
||
parser.add_argument("--gpu", type=str, default="0", help="GPU id，例如 0 或 0,1,2,3")
|
||
parser.add_argument("--epochs_filter", type=str, default="", help="只测试指定 epoch，逗号分隔，如 99,109,119")
|
||
parser.add_argument("--min_epoch", type=int, default=None, help="最小 epoch 过滤")
|
||
parser.add_argument("--max_epoch", type=int, default=None, help="最大 epoch 过滤")
|
||
parser.add_argument("--sleep_sec", type=float, default=1.0, help="每次测试后等待秒数")
|
||
return parser.parse_args()
|
||
|
||
|
||
def load_yaml(path: Path) -> dict:
|
||
with path.open("r", encoding="utf-8") as f:
|
||
return yaml.safe_load(f)
|
||
|
||
|
||
def save_yaml(path: Path, data: dict) -> None:
|
||
with path.open("w", encoding="utf-8") as f:
|
||
yaml.safe_dump(data, f, allow_unicode=True, sort_keys=False)
|
||
|
||
|
||
def list_checkpoints(models_dir: Path) -> List[Tuple[int, Path]]:
|
||
ckpts: List[Tuple[int, Path]] = []
|
||
for p in models_dir.glob("checkpoint_*.pth.tar"):
|
||
m = CKPT_RE.match(p.name)
|
||
if m:
|
||
ckpts.append((int(m.group(1)), p))
|
||
ckpts.sort(key=lambda x: x[0])
|
||
return ckpts
|
||
|
||
|
||
def filter_checkpoints(
|
||
ckpts: List[Tuple[int, Path]],
|
||
epochs_filter: str,
|
||
min_epoch: Optional[int],
|
||
max_epoch: Optional[int],
|
||
) -> List[Tuple[int, Path]]:
|
||
selected = ckpts
|
||
|
||
if epochs_filter.strip():
|
||
wanted = {int(x.strip()) for x in epochs_filter.split(",") if x.strip()}
|
||
selected = [(e, p) for e, p in selected if e in wanted]
|
||
|
||
if min_epoch is not None:
|
||
selected = [(e, p) for e, p in selected if e >= min_epoch]
|
||
|
||
if max_epoch is not None:
|
||
selected = [(e, p) for e, p in selected if e <= max_epoch]
|
||
|
||
return selected
|
||
|
||
|
||
def parse_result_file(path: Path) -> List[Dict]:
|
||
rows: List[Dict] = []
|
||
if not path.exists():
|
||
return rows
|
||
|
||
with path.open("r", encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line or line.startswith("Time"):
|
||
continue
|
||
|
||
m = RESULT_LINE_RE.match(line)
|
||
if not m:
|
||
continue
|
||
|
||
seq = int(m.group(1))
|
||
epoch = int(m.group(2))
|
||
rest = m.group(3).split()
|
||
|
||
if len(rest) < 16:
|
||
continue
|
||
|
||
vals = list(map(float, rest[:16]))
|
||
rows.append(
|
||
{
|
||
"seq": seq,
|
||
"epoch": epoch,
|
||
"AP": vals[0],
|
||
"R100": vals[1],
|
||
"F1": vals[2],
|
||
"R@1": vals[3],
|
||
"R@2": vals[4],
|
||
"R@3": vals[5],
|
||
"R@4": vals[6],
|
||
"R@5": vals[7],
|
||
"R@6": vals[8],
|
||
"R@7": vals[9],
|
||
"R@8": vals[10],
|
||
"R@9": vals[11],
|
||
"R@10": vals[12],
|
||
"R@15": vals[13],
|
||
"R@20": vals[14],
|
||
"R@25": vals[15],
|
||
"raw": line,
|
||
}
|
||
)
|
||
return rows
|
||
|
||
|
||
def make_eval_config(base_config_path: Path, ckpt_path: Path, result_name: str, temp_config_path: Path) -> None:
|
||
cfg = load_yaml(base_config_path)
|
||
exp = cfg["experiment"]
|
||
|
||
exp["train_flag"] = 0
|
||
exp["validate_flag"] = 0
|
||
exp["test_flag"] = 1
|
||
exp["load_model"] = 1
|
||
exp["last_model"] = str(ckpt_path)
|
||
|
||
# 保持原 path_result，不改数据库等输出位置
|
||
save_yaml(temp_config_path, cfg)
|
||
|
||
|
||
def run_one_eval(
|
||
work_dir: Path,
|
||
train_script: Path,
|
||
result_name: str,
|
||
gpu: str,
|
||
tag: str,
|
||
) -> int:
|
||
env = os.environ.copy()
|
||
env["CUDA_VISIBLE_DEVICES"] = gpu
|
||
|
||
cmd = [
|
||
sys.executable,
|
||
str(train_script),
|
||
"--result_name",
|
||
result_name,
|
||
"--gpu",
|
||
gpu,
|
||
"--info",
|
||
f"auto_eval_{tag}",
|
||
]
|
||
|
||
print(f"[INFO][{tag}] Running command: {' '.join(cmd)}")
|
||
print(f"[INFO][{tag}] CUDA_VISIBLE_DEVICES={gpu}")
|
||
print(f"[INFO][{tag}] cwd={work_dir}")
|
||
|
||
proc = subprocess.Popen(
|
||
cmd,
|
||
cwd=str(work_dir),
|
||
env=env,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
bufsize=1,
|
||
universal_newlines=True,
|
||
)
|
||
|
||
try:
|
||
assert proc.stdout is not None
|
||
for line in proc.stdout:
|
||
print(f"[{tag}] {line}", end="")
|
||
proc.wait()
|
||
return proc.returncode
|
||
except KeyboardInterrupt:
|
||
print(f"\n[WARN][{tag}] 收到 Ctrl+C，正在终止当前测试子进程...")
|
||
proc.terminate()
|
||
try:
|
||
proc.wait(timeout=5)
|
||
except Exception:
|
||
print(f"[WARN][{tag}] 子进程未及时退出，强制 kill")
|
||
proc.kill()
|
||
proc.wait()
|
||
raise
|
||
|
||
|
||
def collect_epoch_rows(all_rows: List[Dict], epoch: int) -> List[Dict]:
|
||
return [r for r in all_rows if r["epoch"] == epoch]
|
||
|
||
|
||
def aggregate_rows(rows: List[Dict]) -> Dict[str, float]:
|
||
if not rows:
|
||
return {}
|
||
keys = ["AP", "R100", "F1", "R@1", "R@5", "R@10", "R@25"]
|
||
out = {}
|
||
for k in keys:
|
||
out[f"mean_{k}"] = sum(r[k] for r in rows) / len(rows)
|
||
return out
|
||
|
||
|
||
def save_summary_csv(path: Path, summary: List[Dict]) -> None:
|
||
if not summary:
|
||
return
|
||
fieldnames = list(summary[0].keys())
|
||
with path.open("w", newline="", encoding="utf-8") as f:
|
||
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
||
writer.writeheader()
|
||
writer.writerows(summary)
|
||
|
||
|
||
def run_single_checkpoint(
|
||
epoch: int,
|
||
ckpt: Path,
|
||
gpu: str,
|
||
args: argparse.Namespace,
|
||
project_dir: Path,
|
||
train_script: Path,
|
||
) -> Optional[Dict]:
|
||
tag = f"gpu{gpu}_ep{epoch}"
|
||
temp_root = Path(tempfile.mkdtemp(prefix=f"auto_eval_{tag}_"))
|
||
try:
|
||
# train.py 会优先从 cwd/config.yaml 读取配置
|
||
temp_config = temp_root / "config.yaml"
|
||
make_eval_config(Path(args.config), ckpt, args.result_name, temp_config)
|
||
|
||
# train.py 会把 txt 写到 cwd/result/result_name.txt
|
||
(temp_root / "result").mkdir(parents=True, exist_ok=True)
|
||
|
||
ret = run_one_eval(
|
||
work_dir=temp_root,
|
||
train_script=train_script,
|
||
result_name=args.result_name,
|
||
gpu=gpu,
|
||
tag=tag,
|
||
)
|
||
if ret != 0:
|
||
print(f"[WARN][{tag}] checkpoint {epoch} 测试失败，返回码 {ret}")
|
||
return None
|
||
|
||
time.sleep(args.sleep_sec)
|
||
|
||
result_txt = temp_root / "result" / f"{args.result_name}.txt"
|
||
parsed = parse_result_file(result_txt)
|
||
epoch_rows = collect_epoch_rows(parsed, epoch)
|
||
|
||
if not epoch_rows:
|
||
print(f"[WARN][{tag}] 没有在结果文件中找到 epoch={epoch} 的记录")
|
||
return None
|
||
|
||
agg = aggregate_rows(epoch_rows)
|
||
row = {
|
||
"epoch": epoch,
|
||
"checkpoint": str(ckpt),
|
||
"gpu": gpu,
|
||
**agg,
|
||
}
|
||
print(f"[INFO][{tag}] 汇总: {row}")
|
||
return row
|
||
finally:
|
||
shutil.rmtree(temp_root, ignore_errors=True)
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
|
||
project_dir = Path(args.project_dir).resolve()
|
||
train_script = Path(args.train_script).resolve()
|
||
models_dir = Path(args.models_dir).resolve()
|
||
|
||
if not project_dir.exists():
|
||
raise FileNotFoundError(f"project_dir 不存在: {project_dir}")
|
||
if not Path(args.config).exists():
|
||
raise FileNotFoundError(f"config 不存在: {args.config}")
|
||
if not train_script.exists():
|
||
raise FileNotFoundError(f"train_script 不存在: {train_script}")
|
||
if not models_dir.exists():
|
||
raise FileNotFoundError(f"models_dir 不存在: {models_dir}")
|
||
|
||
ckpts = list_checkpoints(models_dir)
|
||
ckpts = filter_checkpoints(ckpts, args.epochs_filter, args.min_epoch, args.max_epoch)
|
||
|
||
if not ckpts:
|
||
raise RuntimeError("没有找到符合条件的 checkpoint")
|
||
|
||
gpu_list = [x.strip() for x in args.gpu.split(",") if x.strip()]
|
||
if not gpu_list:
|
||
raise RuntimeError("没有可用 GPU 参数")
|
||
|
||
print(f"[INFO] 使用 GPU 列表: {gpu_list}")
|
||
print(f"[INFO] 待测试 checkpoint 数量: {len(ckpts)}")
|
||
|
||
summary_rows: List[Dict] = []
|
||
|
||
# 单卡时保持串行行为
|
||
if len(gpu_list) == 1:
|
||
gpu = gpu_list[0]
|
||
for epoch, ckpt in ckpts:
|
||
print("=" * 100)
|
||
print(f"[INFO] 开始测试 checkpoint: epoch={epoch}, path={ckpt}, gpu={gpu}")
|
||
print("=" * 100)
|
||
row = run_single_checkpoint(epoch, ckpt, gpu, args, project_dir, train_script)
|
||
if row is not None:
|
||
summary_rows.append(row)
|
||
else:
|
||
# 多卡并行：round-robin 分配 checkpoint 到不同 GPU
|
||
futures = []
|
||
with ThreadPoolExecutor(max_workers=len(gpu_list)) as ex:
|
||
for idx, (epoch, ckpt) in enumerate(ckpts):
|
||
gpu = gpu_list[idx % len(gpu_list)]
|
||
futures.append(
|
||
ex.submit(
|
||
run_single_checkpoint,
|
||
epoch,
|
||
ckpt,
|
||
gpu,
|
||
args,
|
||
project_dir,
|
||
train_script,
|
||
)
|
||
)
|
||
|
||
for fut in as_completed(futures):
|
||
row = fut.result()
|
||
if row is not None:
|
||
summary_rows.append(row)
|
||
|
||
summary_rows.sort(key=lambda x: x["epoch"])
|
||
|
||
summary_csv = project_dir / "result" / f"{args.result_name}_summary.csv"
|
||
save_summary_csv(summary_csv, summary_rows)
|
||
print(f"[INFO] 汇总结果已保存到: {summary_csv}")
|
||
|
||
if summary_rows:
|
||
best_by_ap = max(summary_rows, key=lambda x: x.get("mean_AP", float("-inf")))
|
||
print("\n[INFO] 最佳 checkpoint（按 mean_AP）:")
|
||
print(best_by_ap)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |