#!/usr/bin/env python3 """Per-subset evaluation for v12 (ov1_unified_v12). 对 v12 训练得到的 checkpoint（unified_spatial_foa_fsd63_all 数据集上训练） 按子集独立评估，输出每个子集的 F20/ER20/LE_CD/LR_CD/SELD_score 以及 oracle_class_acc / oracle_azi_mae_deg / oracle_ele_mae_deg / oracle_dist_mae 等诊断指标，用于分析： 1) cls vs 空间 谁是瓶颈 2) 真实 vs 仿真 的 gap 3) 多源（ov2/ov3）vs 单源 的差距 支持的子集： - ov1_sim / ov2_sim / ov3_sim （旧仿真 OV 数据，static，所有帧同一 DOA） - ov1_real / ov2_real / ov3_real（DCASE real static mapped，更接近真实混响） - dcase_starss （DCASE STARSS23 / 22 / TAU 混合，real dynamic） - unified_valid / unified_test （新 unified 数据，混合 sim_static + qa_sim + dcase_real） 用法： python eval_v12_per_subset.py \\ --checkpoint checkpoints/spatial_beats_ov1_unified_v12_exp/03_ov123_top4/best.pt \\ --preset ov1_unified_v12 \\ --split valid \\ --batch-size 8 --num-workers 8 --amp bf16 # 跑 test 集合 python eval_v12_per_subset.py \\ --checkpoint checkpoints/spatial_beats_ov1_unified_v12_exp/03_ov123_top4/best.pt \\ --preset ov1_unified_v12 \\ --split test \\ --batch-size 8 --num-workers 8 --amp bf16 """ from __future__ import annotations import argparse import contextlib import copy import dataclasses import functools import json from pathlib import Path from types import SimpleNamespace from typing import Dict, List, Optional, Tuple import torch from tqdm.auto import tqdm from spatial_beats import SpatialBEATs from spatial_dataset import SpatialDataset, collate_spatial_batch from spatial_loss import ( OfficialDCASEMetricsAccumulator, accumulate_frame_track_seld, compute_frame_track_validation_metrics, ) from train_spatial_beats import ( DEFAULT_DCASE_STARSS_VALID_MANIFEST, DEFAULT_OV1_MANIFEST, DEFAULT_OV2_MANIFEST, DEFAULT_OV3_MANIFEST, DEFAULT_OV1_REAL_MANIFEST, DEFAULT_OV2_REAL_MANIFEST, DEFAULT_OV3_REAL_MANIFEST, DEFAULT_UNIFIED_TRAIN_MANIFEST, DEFAULT_UNIFIED_VALID_MANIFEST, TrainSpatialBEATsConfig, build_dataset_config, build_model_config, build_train_config_from_args, ) DEFAULT_DCASE_STARSS_TEST_MANIFEST = ( "/apdcephfs_cq10/share_1603164/user/schmittzhu/data/metadata/" "dcase_starss_foa.test.jsonl" ) DEFAULT_UNIFIED_TEST_MANIFEST = ( "/apdcephfs_cq12/share_302080740/user/schmittzhu/data/" "unified_spatial_foa_fsd63_all/test.jsonl" ) def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument("--checkpoint", required=True) p.add_argument("--preset", required=True) p.add_argument( "--split", choices=("valid", "test"), default="valid", help="valid uses each dataset's valid split; test uses the test split", ) # sim manifests (ov1/2/3) p.add_argument("--ov1-manifest", default=DEFAULT_OV1_MANIFEST) p.add_argument("--ov2-manifest", default=DEFAULT_OV2_MANIFEST) p.add_argument("--ov3-manifest", default=DEFAULT_OV3_MANIFEST) # real (DCASE static mapped) manifests p.add_argument("--ov1-real-manifest", default=DEFAULT_OV1_REAL_MANIFEST) p.add_argument("--ov2-real-manifest", default=DEFAULT_OV2_REAL_MANIFEST) p.add_argument("--ov3-real-manifest", default=DEFAULT_OV3_REAL_MANIFEST) # dcase_starss (full real, dynamic) p.add_argument( "--dcase-starss-valid-manifest", default=DEFAULT_DCASE_STARSS_VALID_MANIFEST, ) p.add_argument( "--dcase-starss-test-manifest", default=DEFAULT_DCASE_STARSS_TEST_MANIFEST, ) # unified p.add_argument("--unified-valid-manifest", default=DEFAULT_UNIFIED_VALID_MANIFEST) p.add_argument("--unified-test-manifest", default=DEFAULT_UNIFIED_TEST_MANIFEST) # runtime p.add_argument("--batch-size", type=int, default=8) p.add_argument("--num-workers", type=int, default=8) p.add_argument("--amp", choices=("fp32", "bf16", "fp16"), default="bf16") p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") p.add_argument("--output-json", default=None) p.add_argument("--activity-threshold", type=float, default=0.5) p.add_argument( "--max-samples-per-subset", type=int, default=-1, help="-1 = no cap (evaluate everything)", ) p.add_argument( "--only-subsets", default="", help="comma-separated subset names to evaluate (empty=all). " "Options: ov1_sim,ov2_sim,ov3_sim,ov1_real,ov2_real,ov3_real," "dcase_starss,unified", ) return p.parse_args() def build_cfg(args: argparse.Namespace) -> TrainSpatialBEATsConfig: ns = SimpleNamespace( preset=args.preset, ov1_manifest=args.ov1_manifest, ov2_manifest=args.ov2_manifest, ov3_manifest=args.ov3_manifest, ov1_real_manifest=args.ov1_real_manifest, ov2_real_manifest=args.ov2_real_manifest, ov3_real_manifest=args.ov3_real_manifest, dcase_starss_valid_manifest=args.dcase_starss_valid_manifest, unified_train_manifest=DEFAULT_UNIFIED_TRAIN_MANIFEST, unified_valid_manifest=args.unified_valid_manifest, batch_size=None, num_workers=None, amp=None, num_epochs=None, learning_rate=None, weight_decay=None, output_dir=None, class_finetuned_ckpt=None, init_from_spatial_ckpt=None, resume=None, no_resume_optimizer=False, reset_epoch_on_resume=False, reset_best_on_resume=False, crop_mode=None, max_clip_duration_seconds=None, save_every_n_epochs=None, train_projector_in_stage1=False, freeze_trunk=False, no_progress=False, distributed=False, local_rank=None, distributed_backend=None, ddp_find_unused_parameters=False, ) cfg = build_train_config_from_args(ns) cfg.batch_size = int(args.batch_size) cfg.num_workers = int(args.num_workers) cfg.amp_dtype = args.amp cfg.distributed = False cfg.show_progress_bars = True cfg.dump_val_predictions = False cfg.num_val_prediction_examples = 0 return cfg def load_model( ckpt_path: str, cfg: TrainSpatialBEATsConfig, device: torch.device ) -> SpatialBEATs: model_cfg = build_model_config(cfg) model = SpatialBEATs(model_cfg) sd = torch.load(ckpt_path, map_location="cpu", weights_only=False) state_dict = sd["model_state_dict"] if "model_state_dict" in sd else sd.get("model", sd) missing, unexpected = model.load_state_dict(state_dict, strict=False) if missing: print(f"[Eval] WARN missing({len(missing)}): {missing[:6]}") if unexpected: print(f"[Eval] WARN unexpected({len(unexpected)}): {unexpected[:6]}") model.to(device).eval() return model def _amp_ctx(dtype: str): if not torch.cuda.is_available(): return contextlib.nullcontext() if dtype == "bf16": return torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16) if dtype == "fp16": return torch.amp.autocast(device_type="cuda", dtype=torch.float16) return contextlib.nullcontext() def _move_to_device(batch, device): field_vals = {} for f in dataclasses.fields(batch): v = getattr(batch, f.name) field_vals[f.name] = v.to(device) if isinstance(v, torch.Tensor) else v return type(batch)(**field_vals) def build_subset_plan(args: argparse.Namespace) -> List[Tuple[str, str, Tuple[str, ...]]]: """返回 [(subset_name, manifest_path, allowed_splits), ...].""" if args.split == "valid": plan = [ ("ov1_sim", args.ov1_manifest, ("valid",)), ("ov2_sim", args.ov2_manifest, ("valid",)), ("ov3_sim", args.ov3_manifest, ("valid",)), ("ov1_real", args.ov1_real_manifest, ("valid",)), ("ov2_real", args.ov2_real_manifest, ("valid",)), ("ov3_real", args.ov3_real_manifest, ("valid",)), ("dcase_starss", args.dcase_starss_valid_manifest, ("valid",)), ("unified", args.unified_valid_manifest, ("valid",)), ] else: # test：旧 ov 数据集有 test split，dcase 也有 plan = [ ("ov1_sim", args.ov1_manifest, ("test",)), ("ov2_sim", args.ov2_manifest, ("test",)), ("ov3_sim", args.ov3_manifest, ("test",)), ("ov1_real", args.ov1_real_manifest, ("test",)), ("ov2_real", args.ov2_real_manifest, ("test",)), ("ov3_real", args.ov3_real_manifest, ("test",)), ("dcase_starss", args.dcase_starss_test_manifest, ("test",)), ("unified", args.unified_test_manifest, ("test",)), ] filter_set = {s.strip() for s in args.only_subsets.split(",") if s.strip()} if filter_set: plan = [entry for entry in plan if entry[0] in filter_set] # drop entries whose manifest file does not exist filtered = [] for name, path, splits in plan: if not Path(path).exists(): print(f"[Eval] SKIP subset={name} (manifest not found: {path})") continue filtered.append((name, path, splits)) return filtered def _empty_running() -> Dict[str, float]: return { "oracle_class_acc": 0.0, "oracle_azi_mae_deg": 0.0, "oracle_ele_mae_deg": 0.0, "oracle_dist_mae": 0.0, "class_acc": 0.0, "azi_mae_deg": 0.0, "ele_mae_deg": 0.0, "dist_mae": 0.0, "activity_precision": 0.0, "activity_recall": 0.0, "activity_acc": 0.0, "matched_count": 0.0, } def eval_one_subset( model: SpatialBEATs, cfg: TrainSpatialBEATsConfig, subset_name: str, manifest_path: str, allowed_splits: Tuple[str, ...], device: torch.device, activity_threshold: float, max_samples: int, ) -> Dict[str, float]: ds_cfg = copy.deepcopy(build_dataset_config(cfg)) ds_cfg.allowed_splits = allowed_splits dataset = SpatialDataset(manifest_path=manifest_path, config=ds_cfg) total = len(dataset) print(f"\n[Eval][{subset_name}] manifest={manifest_path}") print(f"[Eval][{subset_name}] split={allowed_splits} size={total}") if total == 0: print(f"[Eval][{subset_name}] empty — skip") return {"subset": subset_name, "manifest": manifest_path, "size": 0} if max_samples > 0 and total > max_samples: indices = list(range(max_samples)) dataset = torch.utils.data.Subset(dataset, indices) print(f"[Eval][{subset_name}] capped to first {max_samples}") collate = functools.partial(collate_spatial_batch, config=ds_cfg) loader = torch.utils.data.DataLoader( dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.num_workers, collate_fn=collate, pin_memory=True, drop_last=False, persistent_workers=cfg.num_workers > 0, prefetch_factor=4 if cfg.num_workers > 0 else None, ) running = _empty_running() num_batches = 0 seld_acc = OfficialDCASEMetricsAccumulator() with torch.no_grad(): for batch in tqdm(loader, desc=f"Eval {subset_name}", leave=False): batch = _move_to_device(batch, device) with _amp_ctx(cfg.amp_dtype): model_output = model( waveform=batch.waveform, padding_mask=batch.waveform_padding_mask, clip_duration_seconds=batch.clip_duration_seconds, mono_window_mask=None, ) pred_out = model_output.frame_track_prediction_output if pred_out is None: raise RuntimeError("frame_track_prediction_output is None") metric_output = compute_frame_track_validation_metrics( prediction_output=pred_out, batch=batch, temporal_padding_mask=model_output.temporal_padding_mask, config=cfg.loss, ) accumulate_frame_track_seld( prediction_output=pred_out, batch=batch, temporal_padding_mask=model_output.temporal_padding_mask, accumulator=seld_acc, activity_threshold=activity_threshold, ) for key in running: v = getattr(metric_output, key, None) if v is None: continue running[key] += float(v.item()) num_batches += 1 metrics = {k: v / max(num_batches, 1) for k, v in running.items()} metrics.update(seld_acc.compute()) metrics["subset"] = subset_name metrics["manifest"] = manifest_path metrics["size"] = total return metrics def _fmt(v, digits=4): try: return f"{float(v):.{digits}f}" except Exception: return "nan" def print_summary(all_metrics: List[Dict[str, float]], split: str) -> None: if not all_metrics: print("[Eval] no metrics") return print("\n" + "=" * 96) print(f" v12 per-subset ({split} split)") print("=" * 96) hdr = ( f"{'subset':<16} {'N':>6} " f"{'F20':>6} {'ER20':>6} {'LE_CD':>7} {'LR_CD':>6} {'SELD':>6} " f"{'o_cls':>6} {'o_azi':>6} {'o_ele':>6} {'o_dst':>6} " f"{'a_P':>5} {'a_R':>5}" ) print(hdr) print("-" * len(hdr)) for m in all_metrics: if m.get("size", 0) == 0: print(f"{m['subset']:<16} {'EMPTY':>6}") continue print( f"{m['subset']:<16} " f"{m['size']:>6} " f"{_fmt(m.get('F20'), 4):>6} " f"{_fmt(m.get('ER20'), 4):>6} " f"{_fmt(m.get('LE_CD'), 2):>7} " f"{_fmt(m.get('LR_CD'), 4):>6} " f"{_fmt(m.get('SELD_score'), 4):>6} " f"{_fmt(m.get('oracle_class_acc'), 4):>6} " f"{_fmt(m.get('oracle_azi_mae_deg'), 2):>6} " f"{_fmt(m.get('oracle_ele_mae_deg'), 2):>6} " f"{_fmt(m.get('oracle_dist_mae'), 4):>6} " f"{_fmt(m.get('activity_precision'), 3):>5} " f"{_fmt(m.get('activity_recall'), 3):>5}" ) print("=" * 96) print(" legend: F20↑ ER20↓ LE_CD↓ LR_CD↑ SELD↓ " "o_cls=oracle_class_acc o_azi/ele=oracle doa MAE (deg) " "a_P/a_R=activity precision/recall") print("=" * 96) def main() -> None: args = parse_args() device = torch.device(args.device) print(f"[Eval] Device: {device}") print(f"[Eval] Checkpoint: {args.checkpoint}") print(f"[Eval] Preset: {args.preset}") print(f"[Eval] Split: {args.split}") cfg = build_cfg(args) assert cfg.loss.supervision_mode == "local_spatial_track", ( f"Expected local_spatial_track, got {cfg.loss.supervision_mode}." ) if device.type != "cuda": cfg.amp_dtype = "fp32" model = load_model(args.checkpoint, cfg, device) plan = build_subset_plan(args) all_metrics: List[Dict[str, float]] = [] for subset_name, manifest_path, allowed_splits in plan: m = eval_one_subset( model=model, cfg=cfg, subset_name=subset_name, manifest_path=manifest_path, allowed_splits=allowed_splits, device=device, activity_threshold=args.activity_threshold, max_samples=args.max_samples_per_subset, ) all_metrics.append(m) # 打印单子集即时结果 if m.get("size", 0) > 0: print( f"[Eval][{subset_name}] F20={_fmt(m.get('F20'))} " f"ER20={_fmt(m.get('ER20'))} LE_CD={_fmt(m.get('LE_CD'), 2)} " f"LR_CD={_fmt(m.get('LR_CD'))} SELD={_fmt(m.get('SELD_score'))} " f"o_cls={_fmt(m.get('oracle_class_acc'))} " f"o_azi={_fmt(m.get('oracle_azi_mae_deg'), 2)} " f"o_ele={_fmt(m.get('oracle_ele_mae_deg'), 2)} " f"o_dst={_fmt(m.get('oracle_dist_mae'))} " f"a_P={_fmt(m.get('activity_precision'))} " f"a_R={_fmt(m.get('activity_recall'))}" ) print_summary(all_metrics, args.split) out_path = args.output_json if out_path is None: out_path = str( Path(args.checkpoint).parent / f"eval_v12_per_subset_{args.split}.json" ) with open(out_path, "w") as f: json.dump( { "checkpoint": args.checkpoint, "preset": args.preset, "split": args.split, "activity_threshold": args.activity_threshold, "per_subset": all_metrics, }, f, indent=2, ensure_ascii=True, ) print(f"\n[Eval] Summary saved to {out_path}") if __name__ == "__main__": main()