SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-12 20:15:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(benchmark): support multi-dataset evaluation with configurable top-k list

Browse Source 
- 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:
SikongJueluo
2026-05-31 15:56:35 +08:00
parent 7a1e1ccf3f
commit 9eb52f8cef
6 changed files with 454 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

140
mini-nav/benchmarks/tasks/retrieval.py
View File

@@ -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,
}