mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
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:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user