mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
SikongJueluo
ab616528b4
- 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
329 lines
10 KiB
Python
329 lines
10 KiB
Python
"""Retrieval task for benchmark evaluation (Recall@K)."""
|
|
|
|
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.
|
|
|
|
Args:
|
|
vector_dim: Feature vector dimension.
|
|
|
|
Returns:
|
|
PyArrow schema with id, label, and vector fields.
|
|
"""
|
|
return pa.schema(
|
|
[
|
|
pa.field("id", pa.int32()),
|
|
pa.field("label", pa.int32()),
|
|
pa.field("vector", pa.list_(pa.float32(), vector_dim)),
|
|
]
|
|
)
|
|
|
|
|
|
def _establish_eval_database(
|
|
processor: BitImageProcessor,
|
|
model: nn.Module,
|
|
table: lancedb.table.Table,
|
|
dataloader: DataLoader[Any],
|
|
) -> None:
|
|
"""Extract features from training images and store them in a database table.
|
|
|
|
Args:
|
|
processor: Image preprocessor.
|
|
model: Feature extraction model.
|
|
table: LanceDB table to store features.
|
|
dataloader: DataLoader for the training dataset.
|
|
"""
|
|
# 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[config.benchmark.dataset.label_column]
|
|
labels_list = labels.tolist()
|
|
batch_size = len(labels_list)
|
|
|
|
table.add(
|
|
[
|
|
{
|
|
"id": global_idx + j,
|
|
"label": labels_list[j],
|
|
"vector": all_features[global_idx + j].detach().cpu().numpy(),
|
|
}
|
|
for j in range(batch_size)
|
|
]
|
|
)
|
|
global_idx += batch_size
|
|
|
|
|
|
def _evaluate_recall(
|
|
processor: BitImageProcessor,
|
|
model: nn.Module,
|
|
table: lancedb.table.Table,
|
|
dataloader: DataLoader[Any],
|
|
top_k: int,
|
|
) -> tuple[int, int]:
|
|
"""Evaluate Recall@K by searching the database for each test image.
|
|
|
|
Args:
|
|
processor: Image preprocessor.
|
|
model: Feature extraction model.
|
|
table: LanceDB table to search against.
|
|
dataloader: DataLoader for the test dataset.
|
|
top_k: Number of top results to retrieve.
|
|
|
|
Returns:
|
|
A tuple of (correct_count, total_count).
|
|
"""
|
|
# 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[config.benchmark.dataset.label_column]
|
|
labels_list = labels.tolist()
|
|
|
|
for j in range(len(labels_list)):
|
|
feature = all_features[feature_idx + j].tolist()
|
|
true_label = labels_list[j]
|
|
|
|
results = (
|
|
table.search(feature)
|
|
.select(["label", "_distance"])
|
|
.limit(top_k)
|
|
.to_polars()
|
|
)
|
|
|
|
retrieved_labels = results["label"].to_list()
|
|
if true_label in retrieved_labels:
|
|
correct += 1
|
|
total += 1
|
|
|
|
feature_idx += len(labels_list)
|
|
|
|
return correct, total
|
|
|
|
|
|
@RegisterTask("retrieval")
|
|
class RetrievalTask(BaseBenchmarkTask):
|
|
"""Retrieval evaluation task (Recall@K)."""
|
|
|
|
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,
|
|
model: Any,
|
|
processor: Any,
|
|
train_dataset: Any,
|
|
table: lancedb.table.Table,
|
|
batch_size: int,
|
|
) -> None:
|
|
"""Build the evaluation database from training data.
|
|
|
|
Args:
|
|
model: Feature extraction model.
|
|
processor: Image preprocessor.
|
|
train_dataset: Training dataset.
|
|
table: LanceDB table to store features.
|
|
batch_size: Batch size for DataLoader.
|
|
"""
|
|
# Get a sample image to infer vector dimension
|
|
sample = train_dataset[0]
|
|
sample_image = sample["img"]
|
|
|
|
vector_dim = infer_vector_dim(processor, model, sample_image)
|
|
expected_schema = _build_eval_schema(vector_dim)
|
|
|
|
# Check schema compatibility
|
|
if table.schema != expected_schema:
|
|
raise ValueError(
|
|
f"Table schema mismatch. Expected: {expected_schema}, "
|
|
f"Got: {table.schema}"
|
|
)
|
|
|
|
# Build database
|
|
train_loader = DataLoader(
|
|
train_dataset.with_format("torch"),
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
num_workers=4,
|
|
)
|
|
_establish_eval_database(processor, model, table, train_loader)
|
|
|
|
def evaluate(
|
|
self,
|
|
model: Any,
|
|
processor: Any,
|
|
test_dataset: Any,
|
|
table: lancedb.table.Table,
|
|
batch_size: int,
|
|
) -> dict[str, Any]:
|
|
"""Evaluate the model on the test dataset.
|
|
|
|
Args:
|
|
model: Feature extraction model.
|
|
processor: Image preprocessor.
|
|
test_dataset: Test dataset.
|
|
table: LanceDB table to search against.
|
|
batch_size: Batch size for DataLoader.
|
|
|
|
Returns:
|
|
Dictionary containing evaluation results with keys:
|
|
- accuracy: Recall@K accuracy (0.0 ~ 1.0)
|
|
- correct: Number of correct predictions
|
|
- total: Total number of test samples
|
|
- top_k: The K value used
|
|
"""
|
|
test_loader = DataLoader(
|
|
test_dataset.with_format("torch"),
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
num_workers=4,
|
|
)
|
|
correct, total = _evaluate_recall(
|
|
processor, model, table, test_loader, self.top_k
|
|
)
|
|
|
|
accuracy = correct / total if total > 0 else 0.0
|
|
|
|
return {
|
|
"accuracy": accuracy,
|
|
"correct": correct,
|
|
"total": total,
|
|
"top_k": self.top_k,
|
|
}
|