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

refactor(benchmark): delegate model loading to tasks and support CIFAR-100

Browse Source 
- Extract model loading logic from benchmark CLI into task-owned prepare_benchmark
- Add RetrievalEncoder class wrapping DINO with optional hash compression
- Add accelerate dependency for device management  
- Switch dataset from CIFAR-10 to CIFAR-100 with fine_label column
This commit is contained in:
SikongJueluo
2026-05-09 15:12:19 +08:00
parent 0fbcd915bd
commit ab616528b4
7 changed files with 1828 additions and 1226 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
mini-nav/benchmarks/runner.py
View File

@@ -1,18 +1,42 @@
"""Benchmark runner for executing evaluations."""
from pathlib import Path
from typing import Any
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
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": config.task.top_k}
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.
@@ -130,10 +154,30 @@ def _print_benchmark_info(
console.print(table)
def _print_benchmark_results(results: dict[str, Any]) -> None:
"""Print benchmark results using Rich table.
Args:
results: Final benchmark metrics.
"""
table = Table(title="Benchmark Results", show_header=False)
table.add_column("Metric", style="cyan", no_wrap=True)
table.add_column("Value", style="green")
for key, value in results.items():
if isinstance(value, float):
table.add_row(key, f"{value:.4f}")
continue
table.add_row(key, str(value))
console.print(table)
def run_benchmark(
model: Any,
processor: Any,
config: BenchmarkConfig,
model_config: ModelConfig | None = None,
model_name: str = "model",
) -> dict[str, Any]:
"""Run benchmark evaluation.
@@ -148,6 +192,7 @@ def run_benchmark(
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.
Returns:
@@ -171,6 +216,23 @@ def run_benchmark(
f"Dataset {config.dataset.path} does not have train/test splits"
)
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 a sample
sample = train_dataset[0]
sample_image = sample["img"]
@@ -191,16 +253,17 @@ def run_benchmark(
f"[yellow]Table '{table_name}' already has {table_count} entries, skipping database build.[/yellow]"
)
else:
# Create and run benchmark task
task = get_task(config.task.type, top_k=config.task.top_k)
console.print(
f"[cyan]Building database[/cyan] with {len(train_dataset)} training samples..."
)
task.build_database(model, processor, train_dataset, table, config.batch_size)
table_count = table.count_rows()
_print_benchmark_info(config, vector_dim, table_name, table_count)
# Run evaluation (results with Rich table will be printed by the task)
task = get_task(config.task.type, top_k=config.task.top_k)
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)
return results

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

@@ -1,17 +1,62 @@
"""Retrieval task for benchmark evaluation (Recall@K)."""
from typing import Any
from typing import TYPE_CHECKING, Any, cast
import lancedb
import pyarrow as pa
import torch
import torch.nn.functional as F
from benchmarks.base import BaseBenchmarkTask
from benchmarks.tasks.registry import RegisterTask
from compressors.model_loader import get_dino_dim, load_dino_model, load_hash_compressor
from configs import cfg_manager
from rich.progress import track
from torch import nn
from torch.utils.data import DataLoader
from transformers import BitImageProcessor
from utils.feature_extractor import extract_batch_features, infer_vector_dim
if TYPE_CHECKING:
from compressors.hash_compressor import HashCompressor
class RetrievalEncoder(nn.Module):
"""Benchmark encoder for DINO and optional hash compression."""
def __init__(
self,
dino: nn.Module,
compressor: "HashCompressor | None" = None,
) -> None:
"""Initialize retrieval encoder.
Args:
dino: DINO backbone used for feature extraction.
compressor: Optional hash compressor for recall evaluation.
"""
super().__init__()
self.dino: nn.Module = dino
self.compressor: HashCompressor | None = compressor
def forward(self, inputs: Any) -> torch.Tensor:
"""Encode processor inputs into benchmark vectors.
Args:
inputs: Batched processor outputs.
Returns:
Float tensor used for LanceDB insertion and retrieval.
"""
outputs = self.dino(**inputs)
tokens = outputs.last_hidden_state
if self.compressor is None:
features = tokens.mean(dim=1)
return F.normalize(features, dim=-1)
bits = self.compressor.encode(tokens)
return bits.to(dtype=torch.float32)
def _build_eval_schema(vector_dim: int) -> pa.Schema:
"""Build PyArrow schema for evaluation database table.
@@ -35,7 +80,7 @@ def _establish_eval_database(
processor: BitImageProcessor,
model: nn.Module,
table: lancedb.table.Table,
dataloader: DataLoader,
dataloader: DataLoader[Any],
) -> None:
"""Extract features from training images and store them in a database table.
@@ -47,11 +92,12 @@ def _establish_eval_database(
"""
# Extract all features using the utility function
all_features = extract_batch_features(processor, model, dataloader)
config = cfg_manager.get()
# Store features to database
global_idx = 0
for batch in track(dataloader, description="Storing eval database"):
labels = batch["label"]
labels = batch[config.benchmark.dataset.label_column]
labels_list = labels.tolist()
batch_size = len(labels_list)
@@ -72,7 +118,7 @@ def _evaluate_recall(
processor: BitImageProcessor,
model: nn.Module,
table: lancedb.table.Table,
dataloader: DataLoader,
dataloader: DataLoader[Any],
top_k: int,
) -> tuple[int, int]:
"""Evaluate Recall@K by searching the database for each test image.
@@ -89,13 +135,14 @@ def _evaluate_recall(
"""
# Extract all features using the utility function
all_features = extract_batch_features(processor, model, dataloader)
config = cfg_manager.get()
correct = 0
total = 0
feature_idx = 0
for batch in track(dataloader, description=f"Evaluating Recall@{top_k}"):
labels = batch["label"]
labels = batch[config.benchmark.dataset.label_column]
labels_list = labels.tolist()
for j in range(len(labels_list)):
@@ -123,14 +170,79 @@ def _evaluate_recall(
class RetrievalTask(BaseBenchmarkTask):
"""Retrieval evaluation task (Recall@K)."""
def __init__(self, top_k: int = 10):
def __init__(
self,
top_k: int = 10,
dino_model: str = "facebook/dinov2-large",
compression_dim: int = 512,
compressor_path: str | None = None,
):
"""Initialize retrieval task.
Args:
top_k: Number of top results to retrieve for recall calculation.
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
self.dino_model = dino_model
self.compression_dim = compression_dim
self.compressor_path = compressor_path
self._processor: BitImageProcessor | None = None
self._model: nn.Module | None = None
self._model_name = "hash_compressor" if compressor_path else "dinov2"
def prepare_benchmark(
self,
model: Any,
processor: Any,
model_name: str = "model",
) -> tuple[nn.Module, BitImageProcessor, str]:
"""Resolve benchmark resources for this task.
Args:
model: Optional pre-built model from the caller.
processor: Optional pre-built processor from the caller.
model_name: Fallback table model name.
Returns:
Tuple of benchmark model, processor, and resolved model name.
"""
if model is not None and processor is not None:
return (
cast(nn.Module, model),
cast(BitImageProcessor, processor),
model_name,
)
self._ensure_resources_loaded()
return (
cast(nn.Module, self._model),
cast(BitImageProcessor, self._processor),
self._model_name,
)
def _ensure_resources_loaded(self) -> None:
"""Lazy-load retrieval benchmark resources."""
if self._processor is not None and self._model is not None:
return
processor, dino = load_dino_model(self.dino_model)
compressor = None
if self.compressor_path is not None:
compressor = load_hash_compressor(
input_dim=get_dino_dim(self.dino_model),
hash_bits=self.compression_dim,
compressor_path=self.compressor_path,
)
compressor.eval()
self._processor = processor
self._model = RetrievalEncoder(dino=dino, compressor=compressor)
self._model.eval()
def build_database(
self,