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:
2026-05-31 15:56:35 +08:00
parent 7a1e1ccf3f
commit 9eb52f8cef
6 changed files with 454 additions and 96 deletions

View File

@@ -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"

View File

@@ -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(