mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
9eb52f8cef
- Evaluate multiple datasets in a single run (CIFAR10 + CIFAR100) - Report Recall@K for a list of K values from one underlying search - Save results to disk: summary.md, metrics.csv, per-dataset predictions.csv, confusion_matrix.csv - Richer evaluation output: query_labels, topk_ids, topk_labels for downstream analysis - Add --dataset and --top-k CLI overrides for benchmark command - Update config schema: dataset→datasets, top_k→top_k_list
456 lines
15 KiB
Python
456 lines
15 KiB
Python
"""Benchmark runner for executing evaluations."""
|
|
|
|
import csv
|
|
import json
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Callable, cast
|
|
|
|
import lancedb
|
|
from benchmarks.datasets import HuggingFaceDataset, LocalDataset
|
|
from benchmarks.tasks import get_task
|
|
from configs.models import BenchmarkConfig, DatasetSourceConfig, ModelConfig
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
|
|
console = Console()
|
|
|
|
|
|
def _create_task(config: BenchmarkConfig, model_config: ModelConfig | None) -> Any:
|
|
"""Create benchmark task with task-specific model settings.
|
|
|
|
Args:
|
|
config: Benchmark configuration.
|
|
model_config: Optional model configuration for task-owned loading.
|
|
|
|
Returns:
|
|
Benchmark task instance.
|
|
"""
|
|
task_kwargs: dict[str, Any] = {"top_k_list": config.task.top_k_list}
|
|
|
|
if config.task.type == "retrieval" and model_config is not None:
|
|
task_kwargs.update(
|
|
{
|
|
"dino_model": model_config.dino_model,
|
|
"compression_dim": model_config.compression_dim,
|
|
"compressor_path": model_config.compressor_path,
|
|
}
|
|
)
|
|
|
|
return get_task(config.task.type, **task_kwargs)
|
|
|
|
|
|
def create_dataset(config: DatasetSourceConfig) -> Any:
|
|
"""Create a dataset instance from configuration.
|
|
|
|
Args:
|
|
config: Dataset source configuration.
|
|
|
|
Returns:
|
|
Dataset instance.
|
|
|
|
Raises:
|
|
ValueError: If source_type is not supported.
|
|
"""
|
|
if config.source_type == "huggingface":
|
|
return HuggingFaceDataset(
|
|
hf_id=config.path,
|
|
img_column=config.img_column,
|
|
label_column=config.label_column,
|
|
)
|
|
elif config.source_type == "local":
|
|
return LocalDataset(
|
|
local_path=config.path,
|
|
img_column=config.img_column,
|
|
label_column=config.label_column,
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported source_type: {config.source_type}. "
|
|
f"Supported types: 'huggingface', 'local'"
|
|
)
|
|
|
|
|
|
def _get_table_name(
|
|
config: BenchmarkConfig,
|
|
model_name: str,
|
|
dataset_path: str,
|
|
) -> str:
|
|
"""Generate database table name from config, model, and dataset.
|
|
|
|
Args:
|
|
config: Benchmark configuration.
|
|
model_name: Model name for table naming.
|
|
dataset_path: Dataset identifier or path (e.g. "uoft-cs/cifar100").
|
|
|
|
Returns:
|
|
Formatted table name.
|
|
"""
|
|
prefix = config.model_table_prefix
|
|
dataset_name = Path(dataset_path).name.lower().replace("-", "_")
|
|
return f"{prefix}_{dataset_name}_{model_name}"
|
|
|
|
|
|
def _ensure_table(
|
|
config: BenchmarkConfig,
|
|
model_name: str,
|
|
vector_dim: int,
|
|
dataset_path: str,
|
|
) -> lancedb.table.Table:
|
|
"""Ensure the LanceDB table exists with correct schema.
|
|
|
|
Args:
|
|
config: Benchmark configuration.
|
|
model_name: Model name for table naming.
|
|
vector_dim: Feature vector dimension.
|
|
dataset_path: Dataset identifier or path.
|
|
|
|
Returns:
|
|
LanceDB table instance.
|
|
"""
|
|
import pyarrow as pa
|
|
from database import db_manager
|
|
|
|
table_name = _get_table_name(config, model_name, dataset_path)
|
|
|
|
# Build expected schema
|
|
schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.int32()),
|
|
pa.field("label", pa.int32()),
|
|
pa.field("vector", pa.list_(pa.float32(), vector_dim)),
|
|
]
|
|
)
|
|
|
|
db = db_manager.db
|
|
existing_tables = db.list_tables().tables
|
|
|
|
# Check if table exists and has correct schema
|
|
if table_name in existing_tables:
|
|
table = db.open_table(table_name)
|
|
if table.schema != schema:
|
|
console.print(
|
|
f"[yellow]Table '{table_name}' schema mismatch, rebuilding.[/yellow]"
|
|
)
|
|
db.drop_table(table_name)
|
|
table = db.create_table(table_name, schema=schema)
|
|
else:
|
|
table = db.create_table(table_name, schema=schema)
|
|
|
|
return table
|
|
|
|
|
|
def _print_benchmark_info(
|
|
config: BenchmarkConfig,
|
|
vector_dim: int,
|
|
table_name: str,
|
|
table_count: int,
|
|
dataset_path: str = "",
|
|
) -> None:
|
|
"""Print benchmark configuration info using Rich table.
|
|
|
|
Args:
|
|
config: Benchmark configuration.
|
|
vector_dim: Feature vector dimension.
|
|
table_name: Database table name.
|
|
table_count: Number of entries in the table.
|
|
dataset_path: Current dataset identifier or path.
|
|
"""
|
|
table = Table(title="Benchmark Configuration", show_header=False)
|
|
table.add_column("Key", style="cyan", no_wrap=True)
|
|
table.add_column("Value", style="magenta")
|
|
|
|
table.add_row("Dataset", dataset_path)
|
|
table.add_row("Model Output Dimension", str(vector_dim))
|
|
table.add_row("Table Name", table_name)
|
|
table.add_row("Table Entries", str(table_count))
|
|
|
|
console.print(table)
|
|
|
|
|
|
def _print_benchmark_results(results: dict[str, Any]) -> None:
|
|
"""Print benchmark results using Rich table.
|
|
|
|
Args:
|
|
results: Final benchmark metrics for a single dataset.
|
|
"""
|
|
table = Table(title="Benchmark Results", show_header=False)
|
|
table.add_column("Metric", style="cyan", no_wrap=True)
|
|
table.add_column("Value", style="green")
|
|
|
|
# Print recalls first
|
|
recalls = results.get("recalls", {})
|
|
for k_str, v in recalls.items():
|
|
if isinstance(v, float):
|
|
table.add_row(k_str, f"{v:.4f}")
|
|
|
|
for key, value in results.items():
|
|
if key == "recalls":
|
|
continue
|
|
if isinstance(value, (list, dict)):
|
|
continue
|
|
if isinstance(value, float):
|
|
table.add_row(key, f"{value:.4f}")
|
|
continue
|
|
table.add_row(key, str(value))
|
|
|
|
console.print(table)
|
|
|
|
|
|
def _save_benchmark_outputs(
|
|
all_results: dict[str, dict[str, Any]],
|
|
output_root: Path,
|
|
model_name: str,
|
|
) -> Path:
|
|
"""Save multi-dataset benchmark results to disk.
|
|
|
|
Writes to ``output_root/{run_id}/``:
|
|
- summary.md — overall summary with per-dataset recall table
|
|
- metrics.csv — one row per (dataset, K) combination
|
|
- {dataset_name}/predictions.csv — per-sample predictions
|
|
- {dataset_name}/confusion_matrix.csv — row-normalized (Top-1)
|
|
|
|
Args:
|
|
all_results: Dict mapping dataset name to evaluate() result dict.
|
|
output_root: Parent directory (e.g. ``outputs/benchmark``).
|
|
model_name: Model identifier for the run_id.
|
|
|
|
Returns:
|
|
Path to the created run directory.
|
|
"""
|
|
run_id = f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-{model_name}"
|
|
out_dir = output_root / run_id
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# --- metrics.csv: one row per (dataset, k) ---
|
|
csv_keys = ["model", "run_id", "dataset", "k", "recall@k", "total", "num_classes"]
|
|
all_csv_rows: list[dict[str, Any]] = []
|
|
|
|
# --- summary.md lines ---
|
|
md_lines = [
|
|
"# Benchmark Results",
|
|
"",
|
|
f"- **run_id**: `{run_id}`",
|
|
f"- **model**: `{model_name}`",
|
|
"",
|
|
"## Per-Dataset Recall",
|
|
"",
|
|
]
|
|
|
|
for dataset_name, results in all_results.items():
|
|
recalls = results.get("recalls", {})
|
|
total = results.get("total", 0)
|
|
num_classes = results.get("num_classes", 0)
|
|
|
|
md_lines.append(f"### {dataset_name}")
|
|
md_lines.append("")
|
|
md_lines.append(f"- total queries: `{total}`")
|
|
md_lines.append(f"- num_classes: `{num_classes}`")
|
|
md_lines.append("")
|
|
md_lines.append("| K | Recall@K |")
|
|
md_lines.append("|---:|---:|")
|
|
|
|
for k_str, recall_val in sorted(recalls.items()):
|
|
k = int(k_str.replace("recall@", ""))
|
|
md_lines.append(f"| {k} | {recall_val:.4f} |")
|
|
all_csv_rows.append({
|
|
"model": model_name,
|
|
"run_id": run_id,
|
|
"dataset": dataset_name,
|
|
"k": k,
|
|
"recall@k": recall_val,
|
|
"total": total,
|
|
"num_classes": num_classes,
|
|
})
|
|
|
|
md_lines.append("")
|
|
|
|
# Per-dataset subdirectory
|
|
dataset_dir = out_dir / dataset_name
|
|
dataset_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# --- predictions.csv (per-sample) ---
|
|
query_labels = results.get("query_labels", [])
|
|
topk_labels = results.get("topk_labels", [])
|
|
if query_labels and topk_labels:
|
|
max_k = max(len(preds) for preds in topk_labels) if topk_labels else 1
|
|
fieldnames = ["idx", "true_label"] + [
|
|
f"top{i+1}_label" for i in range(max_k)
|
|
]
|
|
with (dataset_dir / "predictions.csv").open(
|
|
"w", newline="", encoding="utf-8"
|
|
) as f:
|
|
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
for idx, (true_label, preds) in enumerate(
|
|
zip(query_labels, topk_labels)
|
|
):
|
|
row_data: dict[str, Any] = {
|
|
"idx": idx,
|
|
"true_label": true_label,
|
|
}
|
|
for i in range(max_k):
|
|
row_data[f"top{i+1}_label"] = (
|
|
preds[i] if i < len(preds) else ""
|
|
)
|
|
writer.writerow(row_data)
|
|
|
|
# --- confusion_matrix.csv ---
|
|
cm = results.get("confusion_matrix")
|
|
if cm is not None:
|
|
with (dataset_dir / "confusion_matrix.csv").open(
|
|
"w", newline="", encoding="utf-8"
|
|
) as fh:
|
|
writer = csv.writer(fh)
|
|
for row in cm:
|
|
writer.writerow([f"{v:.6f}" for v in row])
|
|
|
|
# Write metrics.csv
|
|
if all_csv_rows:
|
|
with (out_dir / "metrics.csv").open("w", newline="", encoding="utf-8") as f:
|
|
writer = csv.DictWriter(f, fieldnames=csv_keys)
|
|
writer.writeheader()
|
|
for row in all_csv_rows:
|
|
writer.writerow(row)
|
|
|
|
# Write summary.md
|
|
(out_dir / "summary.md").write_text("\n".join(md_lines) + "\n", encoding="utf-8")
|
|
|
|
console.print(f"[green]Results saved to:[/green] {out_dir}")
|
|
return out_dir
|
|
|
|
|
|
def run_benchmark(
|
|
model: Any,
|
|
processor: Any,
|
|
config: BenchmarkConfig,
|
|
model_config: ModelConfig | None = None,
|
|
model_name: str = "model",
|
|
output_root: Path | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Run benchmark evaluation across one or more datasets.
|
|
|
|
Workflow:
|
|
1. Load model / processor once
|
|
2. For each dataset in config.datasets:
|
|
a. Create dataset, get splits
|
|
b. Build evaluation database
|
|
c. Evaluate with all K values in top_k_list
|
|
3. Print per-dataset results
|
|
4. Save all results to disk (if ``output_root`` provided)
|
|
|
|
Args:
|
|
model: Feature extraction model.
|
|
processor: Image preprocessor.
|
|
config: Benchmark configuration.
|
|
model_config: Optional model configuration for task-owned loading.
|
|
model_name: Model name for table naming.
|
|
output_root: Optional root directory for saving results.
|
|
|
|
Returns:
|
|
Dict mapping dataset name to evaluation result dict.
|
|
|
|
Raises:
|
|
ValueError: If benchmark is not enabled in config.
|
|
"""
|
|
# ── Resolve model / processor once ──────────────────────────
|
|
task = _create_task(config, model_config)
|
|
|
|
resolver = getattr(task, "prepare_benchmark", None)
|
|
if callable(resolver):
|
|
prepare_benchmark = cast(
|
|
Callable[[Any, Any, str], tuple[Any, Any, str]],
|
|
resolver,
|
|
)
|
|
model, processor, model_name = prepare_benchmark(model, processor, model_name)
|
|
|
|
if model is None or processor is None:
|
|
raise ValueError("Benchmark task did not provide a valid model and processor")
|
|
|
|
# ── Infer vector dimension from first dataset ─────────────────
|
|
first_cfg = config.datasets[0]
|
|
first_dataset = create_dataset(first_cfg)
|
|
first_train = first_dataset.get_train_split()
|
|
sample = first_train[0]
|
|
sample_image = sample["img"]
|
|
|
|
from utils.feature_extractor import infer_vector_dim
|
|
|
|
vector_dim = infer_vector_dim(processor, model, sample_image)
|
|
console.print(f"[cyan]Model output dimension:[/cyan] {vector_dim}")
|
|
|
|
# ── Evaluate each dataset ────────────────────────────────────
|
|
all_results: dict[str, dict[str, Any]] = {}
|
|
|
|
for dataset_cfg in config.datasets:
|
|
console.print(
|
|
f"\n[bold cyan]Dataset:[/bold cyan] "
|
|
f"{dataset_cfg.source_type} - {dataset_cfg.path}"
|
|
)
|
|
|
|
dataset = create_dataset(dataset_cfg)
|
|
train_dataset = dataset.get_train_split()
|
|
test_dataset = dataset.get_test_split()
|
|
|
|
if train_dataset is None or test_dataset is None:
|
|
raise ValueError(
|
|
f"Dataset {dataset_cfg.path} does not have train/test splits"
|
|
)
|
|
|
|
dataset_name = Path(dataset_cfg.path).name.lower()
|
|
table = _ensure_table(config, model_name, vector_dim, dataset_cfg.path)
|
|
table_name = _get_table_name(config, model_name, dataset_cfg.path)
|
|
|
|
# Build database (skip if cached)
|
|
table_count = table.count_rows()
|
|
expected_count = len(train_dataset)
|
|
|
|
if table_count == expected_count:
|
|
console.print(
|
|
f"[yellow]Table '{table_name}' already has {table_count} entries, "
|
|
f"skipping database build.[/yellow]"
|
|
)
|
|
else:
|
|
if table_count > 0:
|
|
console.print(
|
|
f"[yellow]Table '{table_name}' has {table_count} entries "
|
|
f"(expected {expected_count}), rebuilding.[/yellow]"
|
|
)
|
|
# Rebuild: drop and recreate
|
|
from database import db_manager
|
|
db_manager.db.drop_table(table_name)
|
|
table = _ensure_table(config, model_name, vector_dim, dataset_cfg.path)
|
|
|
|
console.print(
|
|
f"[cyan]Building database[/cyan] with {expected_count} training samples..."
|
|
)
|
|
task.build_database(
|
|
model, processor, train_dataset, table, config.batch_size,
|
|
label_column=dataset_cfg.label_column,
|
|
)
|
|
table_count = table.count_rows()
|
|
|
|
_print_benchmark_info(
|
|
config, vector_dim, table_name, table_count,
|
|
dataset_path=dataset_cfg.path,
|
|
)
|
|
|
|
# Evaluate
|
|
console.print(
|
|
f"[cyan]Evaluating[/cyan] on {len(test_dataset)} test samples..."
|
|
)
|
|
results = task.evaluate(
|
|
model, processor, test_dataset, table, config.batch_size,
|
|
label_column=dataset_cfg.label_column,
|
|
)
|
|
all_results[dataset_name] = results
|
|
|
|
# Print per-dataset results
|
|
_print_benchmark_results(results)
|
|
|
|
# ── Save outputs ─────────────────────────────────────────────
|
|
if output_root is not None:
|
|
_save_benchmark_outputs(all_results, output_root, model_name)
|
|
|
|
return all_results
|