mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
feat(benchmark): support multi-dataset evaluation with configurable top-k list
- 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
This commit is contained in:
@@ -1,5 +1,8 @@
|
||||
"""Benchmark runner for executing evaluations."""
|
||||
|
||||
import csv
|
||||
import json
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, cast
|
||||
|
||||
@@ -23,7 +26,7 @@ def _create_task(config: BenchmarkConfig, model_config: ModelConfig | None) -> A
|
||||
Returns:
|
||||
Benchmark task instance.
|
||||
"""
|
||||
task_kwargs: dict[str, Any] = {"top_k": config.task.top_k}
|
||||
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(
|
||||
@@ -68,19 +71,23 @@ def create_dataset(config: DatasetSourceConfig) -> Any:
|
||||
)
|
||||
|
||||
|
||||
def _get_table_name(config: BenchmarkConfig, model_name: str) -> str:
|
||||
"""Generate database table name from config and model name.
|
||||
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
|
||||
# Use dataset path as part of table name (sanitize)
|
||||
dataset_name = Path(config.dataset.path).name.lower().replace("-", "_")
|
||||
dataset_name = Path(dataset_path).name.lower().replace("-", "_")
|
||||
return f"{prefix}_{dataset_name}_{model_name}"
|
||||
|
||||
|
||||
@@ -88,6 +95,7 @@ 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.
|
||||
|
||||
@@ -95,6 +103,7 @@ def _ensure_table(
|
||||
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.
|
||||
@@ -102,7 +111,7 @@ def _ensure_table(
|
||||
import pyarrow as pa
|
||||
from database import db_manager
|
||||
|
||||
table_name = _get_table_name(config, model_name)
|
||||
table_name = _get_table_name(config, model_name, dataset_path)
|
||||
|
||||
# Build expected schema
|
||||
schema = pa.schema(
|
||||
@@ -132,7 +141,11 @@ def _ensure_table(
|
||||
|
||||
|
||||
def _print_benchmark_info(
|
||||
config: BenchmarkConfig, vector_dim: int, table_name: str, table_count: int
|
||||
config: BenchmarkConfig,
|
||||
vector_dim: int,
|
||||
table_name: str,
|
||||
table_count: int,
|
||||
dataset_path: str = "",
|
||||
) -> None:
|
||||
"""Print benchmark configuration info using Rich table.
|
||||
|
||||
@@ -141,12 +154,13 @@ def _print_benchmark_info(
|
||||
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", f"{config.dataset.source_type} - {config.dataset.path}")
|
||||
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))
|
||||
@@ -158,13 +172,23 @@ def _print_benchmark_results(results: dict[str, Any]) -> None:
|
||||
"""Print benchmark results using Rich table.
|
||||
|
||||
Args:
|
||||
results: Final benchmark metrics.
|
||||
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
|
||||
@@ -173,20 +197,147 @@ def _print_benchmark_results(results: dict[str, Any]) -> None:
|
||||
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.
|
||||
"""Run benchmark evaluation across one or more datasets.
|
||||
|
||||
Workflow:
|
||||
1. Create dataset from configuration
|
||||
2. Create benchmark task from configuration
|
||||
3. Build evaluation database from training set
|
||||
4. Evaluate on test set
|
||||
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.
|
||||
@@ -194,28 +345,15 @@ def run_benchmark(
|
||||
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:
|
||||
Dictionary containing evaluation results.
|
||||
Dict mapping dataset name to evaluation result dict.
|
||||
|
||||
Raises:
|
||||
ValueError: If benchmark is not enabled in config.
|
||||
"""
|
||||
# Create dataset
|
||||
console.print(
|
||||
f"[cyan]Loading dataset:[/cyan] {config.dataset.source_type} - {config.dataset.path}"
|
||||
)
|
||||
dataset = create_dataset(config.dataset)
|
||||
|
||||
# Get train and test splits
|
||||
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 {config.dataset.path} does not have train/test splits"
|
||||
)
|
||||
|
||||
# ── Resolve model / processor once ──────────────────────────
|
||||
task = _create_task(config, model_config)
|
||||
|
||||
resolver = getattr(task, "prepare_benchmark", None)
|
||||
@@ -224,17 +362,16 @@ def run_benchmark(
|
||||
Callable[[Any, Any, str], tuple[Any, Any, str]],
|
||||
resolver,
|
||||
)
|
||||
model, processor, model_name = prepare_benchmark(
|
||||
model,
|
||||
processor,
|
||||
model_name,
|
||||
)
|
||||
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 a sample
|
||||
sample = train_dataset[0]
|
||||
# ── 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
|
||||
@@ -242,28 +379,77 @@ def run_benchmark(
|
||||
vector_dim = infer_vector_dim(processor, model, sample_image)
|
||||
console.print(f"[cyan]Model output dimension:[/cyan] {vector_dim}")
|
||||
|
||||
# Ensure table exists with correct schema
|
||||
table = _ensure_table(config, model_name, vector_dim)
|
||||
table_name = _get_table_name(config, model_name)
|
||||
# ── Evaluate each dataset ────────────────────────────────────
|
||||
all_results: dict[str, dict[str, Any]] = {}
|
||||
|
||||
# Check if database is already built
|
||||
table_count = table.count_rows()
|
||||
if table_count > 0:
|
||||
for dataset_cfg in config.datasets:
|
||||
console.print(
|
||||
f"[yellow]Table '{table_name}' already has {table_count} entries, skipping database build.[/yellow]"
|
||||
f"\n[bold cyan]Dataset:[/bold cyan] "
|
||||
f"{dataset_cfg.source_type} - {dataset_cfg.path}"
|
||||
)
|
||||
else:
|
||||
console.print(
|
||||
f"[cyan]Building database[/cyan] with {len(train_dataset)} training samples..."
|
||||
)
|
||||
task.build_database(model, processor, train_dataset, table, config.batch_size)
|
||||
|
||||
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)
|
||||
|
||||
_print_benchmark_info(config, vector_dim, table_name, table_count)
|
||||
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)
|
||||
|
||||
# Run evaluation (results with Rich table will be printed by the task)
|
||||
console.print(f"[cyan]Evaluating[/cyan] on {len(test_dataset)} test samples...")
|
||||
results = task.evaluate(model, processor, test_dataset, table, config.batch_size)
|
||||
_print_benchmark_results(results)
|
||||
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()
|
||||
|
||||
return results
|
||||
_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
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
import lancedb
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
@@ -81,6 +82,7 @@ def _establish_eval_database(
|
||||
model: nn.Module,
|
||||
table: lancedb.table.Table,
|
||||
dataloader: DataLoader[Any],
|
||||
label_column: str = "label",
|
||||
) -> None:
|
||||
"""Extract features from training images and store them in a database table.
|
||||
|
||||
@@ -89,6 +91,7 @@ def _establish_eval_database(
|
||||
model: Feature extraction model.
|
||||
table: LanceDB table to store features.
|
||||
dataloader: DataLoader for the training dataset.
|
||||
label_column: Column name for labels in the batch dict.
|
||||
"""
|
||||
# Extract all features using the utility function
|
||||
all_features = extract_batch_features(processor, model, dataloader)
|
||||
@@ -97,7 +100,7 @@ def _establish_eval_database(
|
||||
# Store features to database
|
||||
global_idx = 0
|
||||
for batch in track(dataloader, description="Storing eval database"):
|
||||
labels = batch[config.benchmark.dataset.label_column]
|
||||
labels = batch[label_column]
|
||||
labels_list = labels.tolist()
|
||||
batch_size = len(labels_list)
|
||||
|
||||
@@ -120,7 +123,8 @@ def _evaluate_recall(
|
||||
table: lancedb.table.Table,
|
||||
dataloader: DataLoader[Any],
|
||||
top_k: int,
|
||||
) -> tuple[int, int]:
|
||||
label_column: str = "label",
|
||||
) -> dict[str, Any]:
|
||||
"""Evaluate Recall@K by searching the database for each test image.
|
||||
|
||||
Args:
|
||||
@@ -129,9 +133,15 @@ def _evaluate_recall(
|
||||
table: LanceDB table to search against.
|
||||
dataloader: DataLoader for the test dataset.
|
||||
top_k: Number of top results to retrieve.
|
||||
label_column: Column name for labels in the batch dict.
|
||||
|
||||
Returns:
|
||||
A tuple of (correct_count, total_count).
|
||||
A dict with keys:
|
||||
- correct: Number of correct predictions
|
||||
- total: Total number of test samples
|
||||
- query_labels: True labels for each query (length N)
|
||||
- topk_ids: IDs of top-K retrieved items (N x K)
|
||||
- topk_labels: Labels of top-K retrieved items (N x K)
|
||||
"""
|
||||
# Extract all features using the utility function
|
||||
all_features = extract_batch_features(processor, model, dataloader)
|
||||
@@ -140,9 +150,12 @@ def _evaluate_recall(
|
||||
correct = 0
|
||||
total = 0
|
||||
feature_idx = 0
|
||||
query_labels: list[int] = []
|
||||
topk_ids: list[list[int]] = []
|
||||
topk_labels: list[list[int]] = []
|
||||
|
||||
for batch in track(dataloader, description=f"Evaluating Recall@{top_k}"):
|
||||
labels = batch[config.benchmark.dataset.label_column]
|
||||
labels = batch[label_column]
|
||||
labels_list = labels.tolist()
|
||||
|
||||
for j in range(len(labels_list)):
|
||||
@@ -151,19 +164,67 @@ def _evaluate_recall(
|
||||
|
||||
results = (
|
||||
table.search(feature)
|
||||
.select(["label", "_distance"])
|
||||
.select(["id", "label", "_distance"])
|
||||
.limit(top_k)
|
||||
.to_polars()
|
||||
)
|
||||
|
||||
retrieved_ids = results["id"].to_list()
|
||||
retrieved_labels = results["label"].to_list()
|
||||
if true_label in retrieved_labels:
|
||||
correct += 1
|
||||
total += 1
|
||||
|
||||
query_labels.append(true_label)
|
||||
topk_ids.append(retrieved_ids)
|
||||
topk_labels.append(retrieved_labels)
|
||||
|
||||
feature_idx += len(labels_list)
|
||||
|
||||
return correct, total
|
||||
return {
|
||||
"correct": correct,
|
||||
"total": total,
|
||||
"query_labels": query_labels,
|
||||
"topk_ids": topk_ids,
|
||||
"topk_labels": topk_labels,
|
||||
}
|
||||
|
||||
|
||||
def _build_confusion_matrix(
|
||||
query_labels: list[int],
|
||||
topk_labels: list[list[int]],
|
||||
num_classes: int,
|
||||
) -> np.ndarray:
|
||||
"""Build row-normalized confusion matrix from Top-1 predictions.
|
||||
|
||||
For each query, the true label is compared against the first
|
||||
retrieved label (Top-1). The resulting matrix is row-normalized
|
||||
so that each row sums to 1.0, representing "given true class X,
|
||||
what fraction of queries were predicted as each class Y".
|
||||
|
||||
Args:
|
||||
query_labels: True labels for each query (length N).
|
||||
topk_labels: Top-K retrieved labels for each query (N x K).
|
||||
num_classes: Total number of label classes.
|
||||
|
||||
Returns:
|
||||
Row-normalized confusion matrix as (num_classes, num_classes)
|
||||
float64 array. Rows with zero queries remain all-zero.
|
||||
"""
|
||||
cm = np.zeros((num_classes, num_classes), dtype=np.int64)
|
||||
for y_true, top_labels in zip(query_labels, topk_labels):
|
||||
y_pred = top_labels[0] # Top-1
|
||||
cm[int(y_true), int(y_pred)] += 1
|
||||
|
||||
# Row normalize: each row sums to 1.0
|
||||
row_sums = cm.sum(axis=1, keepdims=True)
|
||||
cm_norm = np.divide(
|
||||
cm.astype(np.float64),
|
||||
row_sums,
|
||||
out=np.zeros_like(cm, dtype=np.float64),
|
||||
where=row_sums > 0,
|
||||
)
|
||||
return cm_norm
|
||||
|
||||
|
||||
@RegisterTask("retrieval")
|
||||
@@ -173,6 +234,7 @@ class RetrievalTask(BaseBenchmarkTask):
|
||||
def __init__(
|
||||
self,
|
||||
top_k: int = 10,
|
||||
top_k_list: list[int] | None = None,
|
||||
dino_model: str = "facebook/dinov2-large",
|
||||
compression_dim: int = 512,
|
||||
compressor_path: str | None = None,
|
||||
@@ -180,13 +242,22 @@ class RetrievalTask(BaseBenchmarkTask):
|
||||
"""Initialize retrieval task.
|
||||
|
||||
Args:
|
||||
top_k: Number of top results to retrieve for recall calculation.
|
||||
top_k: Maximum K for retrieval search.
|
||||
Deprecated — prefer ``top_k_list``.
|
||||
top_k_list: List of K values to evaluate (all derived from
|
||||
a single max-K search). When not provided, ``[top_k]``
|
||||
is used as fallback.
|
||||
dino_model: DINO model name used for feature extraction.
|
||||
compression_dim: Output dimension of the hash compressor.
|
||||
compressor_path: Optional path to trained hash compressor weights.
|
||||
"""
|
||||
super().__init__(top_k=top_k)
|
||||
self.top_k = top_k
|
||||
if top_k_list is not None:
|
||||
top_k_list = sorted(set(top_k_list))
|
||||
else:
|
||||
top_k_list = [top_k]
|
||||
super().__init__(top_k_list=top_k_list)
|
||||
self.top_k_list: list[int] = top_k_list
|
||||
self.max_k: int = max(self.top_k_list)
|
||||
self.dino_model = dino_model
|
||||
self.compression_dim = compression_dim
|
||||
self.compressor_path = compressor_path
|
||||
@@ -251,6 +322,7 @@ class RetrievalTask(BaseBenchmarkTask):
|
||||
train_dataset: Any,
|
||||
table: lancedb.table.Table,
|
||||
batch_size: int,
|
||||
label_column: str = "label",
|
||||
) -> None:
|
||||
"""Build the evaluation database from training data.
|
||||
|
||||
@@ -260,6 +332,7 @@ class RetrievalTask(BaseBenchmarkTask):
|
||||
train_dataset: Training dataset.
|
||||
table: LanceDB table to store features.
|
||||
batch_size: Batch size for DataLoader.
|
||||
label_column: Column name for labels in the batch dict.
|
||||
"""
|
||||
# Get a sample image to infer vector dimension
|
||||
sample = train_dataset[0]
|
||||
@@ -282,7 +355,7 @@ class RetrievalTask(BaseBenchmarkTask):
|
||||
shuffle=False,
|
||||
num_workers=4,
|
||||
)
|
||||
_establish_eval_database(processor, model, table, train_loader)
|
||||
_establish_eval_database(processor, model, table, train_loader, label_column)
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
@@ -291,6 +364,7 @@ class RetrievalTask(BaseBenchmarkTask):
|
||||
test_dataset: Any,
|
||||
table: lancedb.table.Table,
|
||||
batch_size: int,
|
||||
label_column: str = "label",
|
||||
) -> dict[str, Any]:
|
||||
"""Evaluate the model on the test dataset.
|
||||
|
||||
@@ -300,13 +374,17 @@ class RetrievalTask(BaseBenchmarkTask):
|
||||
test_dataset: Test dataset.
|
||||
table: LanceDB table to search against.
|
||||
batch_size: Batch size for DataLoader.
|
||||
label_column: Column name for labels in the batch dict.
|
||||
|
||||
Returns:
|
||||
Dictionary containing evaluation results with keys:
|
||||
- accuracy: Recall@K accuracy (0.0 ~ 1.0)
|
||||
- correct: Number of correct predictions
|
||||
- recalls: Dict of ``{"recall@K": value}`` for each K
|
||||
- total: Total number of test samples
|
||||
- top_k: The K value used
|
||||
- top_k_list: List of K values evaluated
|
||||
- num_classes: Number of label classes
|
||||
- confusion_matrix: Row-normalized confusion matrix (Top-1)
|
||||
- query_labels: True labels for each query sample
|
||||
- topk_labels: Top-K retrieved labels per query sample
|
||||
"""
|
||||
test_loader = DataLoader(
|
||||
test_dataset.with_format("torch"),
|
||||
@@ -314,15 +392,39 @@ class RetrievalTask(BaseBenchmarkTask):
|
||||
shuffle=False,
|
||||
num_workers=4,
|
||||
)
|
||||
correct, total = _evaluate_recall(
|
||||
processor, model, table, test_loader, self.top_k
|
||||
eval_data = _evaluate_recall(
|
||||
processor, model, table, test_loader, self.max_k, label_column
|
||||
)
|
||||
|
||||
accuracy = correct / total if total > 0 else 0.0
|
||||
total = eval_data["total"]
|
||||
query_labels = eval_data["query_labels"]
|
||||
topk_labels = eval_data["topk_labels"]
|
||||
|
||||
# Compute Recall@k for each k in top_k_list
|
||||
recalls: dict[str, float] = {}
|
||||
for k in self.top_k_list:
|
||||
hits = sum(
|
||||
1
|
||||
for true, preds in zip(query_labels, topk_labels)
|
||||
if true in preds[:k]
|
||||
)
|
||||
recalls[f"recall@{k}"] = hits / total if total > 0 else 0.0
|
||||
|
||||
# Infer number of classes from the data
|
||||
all_labels = query_labels + [
|
||||
label for labels in topk_labels for label in labels
|
||||
]
|
||||
num_classes = max(all_labels) + 1 if all_labels else 1
|
||||
|
||||
# Confusion matrix from Top-1 only
|
||||
cm = _build_confusion_matrix(query_labels, topk_labels, num_classes)
|
||||
|
||||
return {
|
||||
"accuracy": accuracy,
|
||||
"correct": correct,
|
||||
"recalls": recalls,
|
||||
"total": total,
|
||||
"top_k": self.top_k,
|
||||
"top_k_list": self.top_k_list,
|
||||
"num_classes": num_classes,
|
||||
"confusion_matrix": cm.tolist(),
|
||||
"query_labels": query_labels,
|
||||
"topk_labels": topk_labels,
|
||||
}
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
"""Benchmark CLI command."""
|
||||
|
||||
import typer
|
||||
from commands import app
|
||||
|
||||
@@ -8,19 +10,59 @@ def benchmark(
|
||||
model_path: str | None = typer.Option(
|
||||
None, "--model", "-m", help="Path to compressor model weights"
|
||||
),
|
||||
dataset: list[str] | None = typer.Option(
|
||||
None,
|
||||
"--dataset",
|
||||
"-d",
|
||||
help=(
|
||||
"Override datasets (HuggingFace path, repeatable). "
|
||||
"Each value creates a single-entry dataset list."
|
||||
),
|
||||
),
|
||||
top_k: list[int] | None = typer.Option(
|
||||
None,
|
||||
"--top-k",
|
||||
"-k",
|
||||
help="Override top-K values (repeatable). E.g. --top-k 1 --top-k 5 --top-k 10",
|
||||
),
|
||||
):
|
||||
"""Run benchmark evaluation across configured datasets.
|
||||
|
||||
Without flags, runs all datasets and K values from config.yaml.
|
||||
Use --dataset to override datasets, --top-k to override K values.
|
||||
"""
|
||||
from benchmarks import run_benchmark
|
||||
from configs import cfg_manager
|
||||
from configs.models import DatasetSourceConfig
|
||||
|
||||
config = cfg_manager.get()
|
||||
benchmark_cfg = config.benchmark
|
||||
|
||||
if model_path:
|
||||
config.model.compressor_path = model_path
|
||||
|
||||
_ = run_benchmark(
|
||||
# CLI overrides
|
||||
if dataset:
|
||||
benchmark_cfg.datasets = [
|
||||
DatasetSourceConfig(
|
||||
source_type="huggingface",
|
||||
path=d,
|
||||
)
|
||||
for d in dataset
|
||||
]
|
||||
|
||||
if top_k:
|
||||
benchmark_cfg.task.top_k_list = sorted(set(top_k))
|
||||
|
||||
model_name = "hash_compressor" if config.model.compressor_path else "dinov2"
|
||||
output_root = config.output.directory / "benchmark"
|
||||
|
||||
results = run_benchmark(
|
||||
model=None,
|
||||
processor=None,
|
||||
config=benchmark_cfg,
|
||||
model_config=config.model,
|
||||
model_name="hash_compressor" if config.model.compressor_path else "dinov2",
|
||||
model_name=model_name,
|
||||
output_root=output_root,
|
||||
)
|
||||
return results
|
||||
|
||||
@@ -22,14 +22,18 @@ dataset:
|
||||
seed: 42
|
||||
|
||||
benchmark:
|
||||
dataset:
|
||||
source_type: "huggingface"
|
||||
path: "uoft-cs/cifar100"
|
||||
img_column: "img"
|
||||
label_column: "fine_label"
|
||||
datasets:
|
||||
- source_type: "huggingface"
|
||||
path: "uoft-cs/cifar100"
|
||||
img_column: "img"
|
||||
label_column: "fine_label"
|
||||
- source_type: "huggingface"
|
||||
path: "uoft-cs/cifar10"
|
||||
img_column: "img"
|
||||
label_column: "label"
|
||||
task:
|
||||
name: "recall_at_k"
|
||||
type: "retrieval"
|
||||
top_k: 1
|
||||
top_k_list: [1, 5, 10]
|
||||
batch_size: 64
|
||||
model_table_prefix: "benchmark"
|
||||
|
||||
@@ -129,7 +129,24 @@ class BenchmarkTaskConfig(BaseModel):
|
||||
|
||||
name: str = Field(default="recall_at_k", description="Task name")
|
||||
type: str = Field(default="retrieval", description="Task type")
|
||||
top_k: int = Field(default=10, gt=0, description="Top K for recall evaluation")
|
||||
top_k_list: list[int] = Field(
|
||||
default=[1, 5, 10],
|
||||
description="Top-K values to evaluate (all derived from a single max-K search)",
|
||||
)
|
||||
|
||||
@property
|
||||
def max_k(self) -> int:
|
||||
"""Maximum K for the underlying search; all values in top_k_list <= max_k."""
|
||||
return max(self.top_k_list) if self.top_k_list else 1
|
||||
|
||||
@field_validator("top_k_list", mode="after")
|
||||
@classmethod
|
||||
def validate_top_k_list(cls, v: list[int]) -> list[int]:
|
||||
if not v:
|
||||
raise ValueError("top_k_list must contain at least one value")
|
||||
if any(k <= 0 for k in v):
|
||||
raise ValueError("top_k_list values must be positive")
|
||||
return sorted(set(v))
|
||||
|
||||
# Multi-object retrieval specific settings
|
||||
gamma: float = Field(
|
||||
@@ -148,7 +165,10 @@ class BenchmarkConfig(BaseModel):
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
dataset: DatasetSourceConfig = Field(default_factory=DatasetSourceConfig)
|
||||
datasets: list[DatasetSourceConfig] = Field(
|
||||
default_factory=lambda: [DatasetSourceConfig()],
|
||||
description="Dataset configurations to evaluate (supports multiple).",
|
||||
)
|
||||
task: BenchmarkTaskConfig = Field(default_factory=BenchmarkTaskConfig)
|
||||
batch_size: int = Field(default=64, gt=0, description="Batch size for DataLoader")
|
||||
model_table_prefix: str = Field(
|
||||
|
||||
Reference in New Issue
Block a user