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

- 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:
2026-05-09 15:12:19 +08:00
parent 0fbcd915bd
commit ab616528b4
7 changed files with 1828 additions and 1226 deletions

View File

@@ -3,11 +3,13 @@
from typing import Any, List, Union, cast
import torch
from aiohttp.web import get
from PIL import Image
from rich.progress import track
from torch import nn
from torch.utils.data import DataLoader
from transformers import BitImageProcessor
from rich.progress import track
from utils import get_device
def _extract_features_from_output(output: Any) -> torch.Tensor:
@@ -26,6 +28,7 @@ def _extract_features_from_output(output: Any) -> torch.Tensor:
return cast(torch.Tensor, output)
@torch.no_grad()
def infer_vector_dim(
processor: BitImageProcessor,
model: nn.Module,
@@ -41,13 +44,9 @@ def infer_vector_dim(
Returns:
Vector dimension.
"""
device = next(model.parameters()).device
model.eval()
with torch.no_grad():
inputs = processor(images=sample_image, return_tensors="pt")
inputs.to(device)
output = model(inputs)
inputs = processor(images=sample_image, return_tensors="pt")
inputs.to(get_device())
output = model(inputs)
features = _extract_features_from_output(output)
return features.shape[-1]
@@ -69,11 +68,8 @@ def extract_single_image_feature(
Returns:
The extracted CLS token feature vector as a list of floats.
"""
device = next(model.parameters()).device
model.eval()
inputs = processor(images=image, return_tensors="pt")
inputs.to(device, non_blocking=True)
inputs.to(get_device(), non_blocking=True)
outputs = model(inputs)
features = _extract_features_from_output(outputs) # [1, D]
@@ -98,8 +94,7 @@ def extract_batch_features(
Returns:
Tensor of shape [batch_size, feature_dim].
"""
device = next(model.parameters()).device
model.eval()
device = get_device()
# Handle DataLoader input
if isinstance(images, DataLoader):