diff --git a/.justfile b/.justfile index 6f8e0be..9543111 100644 --- a/.justfile +++ b/.justfile @@ -25,8 +25,12 @@ ssh: -L 127.0.0.1:2718:172.30.0.2:2718 \ {{ remote_ssh_target }} -docker: - docker exec -it {{ remote_docker_container }} bash +docker cmd="": + @if [ -z "{{ cmd }}" ]; then \ + docker exec -it -w {{ env("REMOTE_WORKDIR") }} {{ remote_docker_container }} bash; \ + else \ + docker exec -i -w {{ env("REMOTE_WORKDIR") }} {{ remote_docker_container }} bash -lc {{ quote(cmd) }}; \ + fi marimo +notebook: uv run marimo edit {{ notebook }} --host 0.0.0.0 --port 2718 --no-token @@ -84,36 +88,36 @@ cam-test-retrieval-write-noise: # Prepare CIFAR10 hash artifact for CAM retrieval smoke benchmark cam-prepare-retrieval-cifar10 ROWS="512" QUERIES="128": - just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar10 --num-rows {{ROWS}} --max-queries {{QUERIES}} --compressor-path outputs/hash_compressor.pt" + just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar10 --num-rows {{ ROWS }} --max-queries {{ QUERIES }} --compressor-path outputs/hash_compressor.pt" # Prepare CIFAR100 hash artifact for CAM retrieval benchmark cam-prepare-retrieval-cifar100 ROWS="512" QUERIES="128": - just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar100 --num-rows {{ROWS}} --max-queries {{QUERIES}} --compressor-path outputs/hash_compressor.pt" + just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar100 --num-rows {{ ROWS }} --max-queries {{ QUERIES }} --compressor-path outputs/hash_compressor.pt" # Run CAM retrieval benchmark on a prepared artifact without hardware noise cam-test-retrieval-artifact DATASET_PATH NUM_ROWS="4096": - just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{NUM_ROWS}} WRITE_NOISE_EN=0 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}" + just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{ NUM_ROWS }} WRITE_NOISE_EN=0 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}" # Run CAM retrieval benchmark on a prepared artifact with write noise enabled (Phase 2: read noise removed) cam-test-retrieval-artifact-write-noise DATASET_PATH NUM_ROWS="4096": - just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{NUM_ROWS}} WRITE_NOISE_EN=1 WRITE_NOISE_RATE_NUM=1 WRITE_NOISE_RATE_DEN=100 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}" + just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{ NUM_ROWS }} WRITE_NOISE_EN=1 WRITE_NOISE_RATE_NUM=1 WRITE_NOISE_RATE_DEN=100 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}" # ── CAM retrieval benchmark noise sweep ──────────────────────────────────────── - # Run noise sweep on a prepared dataset (0%–100%, step 10%) + # Usage: just cam-benchmark-retrieval-sweep DATASET=outputs/.../cifar10_hash512_rows512_queries128.npz NUM_ROWS=512 cam-benchmark-retrieval-sweep DATASET NUM_ROWS="512": - just remote "python scripts/run_retrieval_noise_sweep.py --dataset {{DATASET}} --num-rows {{NUM_ROWS}} --output docs/cam_retrieval_noise_sweep.md" + just remote "python scripts/run_retrieval_noise_sweep.py --dataset {{ DATASET }} --num-rows {{ NUM_ROWS }} --output docs/cam_retrieval_noise_sweep.md" # Prepare CIFAR10 dataset + run full noise sweep (all-in-one) cam-benchmark-sweep-cifar10 ROWS="512" QUERIES="128": - just cam-prepare-retrieval-cifar10 {{ROWS}} {{QUERIES}} - just remote "python scripts/run_retrieval_noise_sweep.py --dataset outputs/cam_retrieval_benchmark/datasets/cifar10_hash512_rows{{ROWS}}_queries{{QUERIES}}.npz --num-rows {{ROWS}} --output docs/cam_retrieval_noise_sweep_cifar10.md" + just cam-prepare-retrieval-cifar10 {{ ROWS }} {{ QUERIES }} + just remote "python scripts/run_retrieval_noise_sweep.py --dataset outputs/cam_retrieval_benchmark/datasets/cifar10_hash512_rows{{ ROWS }}_queries{{ QUERIES }}.npz --num-rows {{ ROWS }} --output docs/cam_retrieval_noise_sweep_cifar10.md" # Prepare CIFAR100 dataset + run full noise sweep (all-in-one) cam-benchmark-sweep-cifar100 ROWS="512" QUERIES="128": - just cam-prepare-retrieval-cifar100 {{ROWS}} {{QUERIES}} - just remote "python scripts/run_retrieval_noise_sweep.py --dataset outputs/cam_retrieval_benchmark/datasets/cifar100_hash512_rows{{ROWS}}_queries{{QUERIES}}.npz --num-rows {{ROWS}} --output docs/cam_retrieval_noise_sweep_cifar100.md" + just cam-prepare-retrieval-cifar100 {{ ROWS }} {{ QUERIES }} + just remote "python scripts/run_retrieval_noise_sweep.py --dataset outputs/cam_retrieval_benchmark/datasets/cifar100_hash512_rows{{ ROWS }}_queries{{ QUERIES }}.npz --num-rows {{ ROWS }} --output docs/cam_retrieval_noise_sweep_cifar100.md" # ── Remote ↔ local sync ──────────────────────────────────────────────────────── diff --git a/mini-nav/benchmarks/runner.py b/mini-nav/benchmarks/runner.py index 1b20aa1..ef56054 100644 --- a/mini-nav/benchmarks/runner.py +++ b/mini-nav/benchmarks/runner.py @@ -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 diff --git a/mini-nav/benchmarks/tasks/retrieval.py b/mini-nav/benchmarks/tasks/retrieval.py index 2deae17..c0142b0 100644 --- a/mini-nav/benchmarks/tasks/retrieval.py +++ b/mini-nav/benchmarks/tasks/retrieval.py @@ -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, } diff --git a/mini-nav/commands/benchmark.py b/mini-nav/commands/benchmark.py index ecca847..7ea3e9a 100644 --- a/mini-nav/commands/benchmark.py +++ b/mini-nav/commands/benchmark.py @@ -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 diff --git a/mini-nav/configs/config.yaml b/mini-nav/configs/config.yaml index 220e63b..1a028f7 100644 --- a/mini-nav/configs/config.yaml +++ b/mini-nav/configs/config.yaml @@ -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" diff --git a/mini-nav/configs/models.py b/mini-nav/configs/models.py index 8df4e04..950e351 100644 --- a/mini-nav/configs/models.py +++ b/mini-nav/configs/models.py @@ -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(