mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
b5a40819cc
- Replace scalar hamming distance with NumPy bitwise_count for batch retrieval - Add ThreadPoolExecutor-based multi-worker query parallelism - Improve missing dataset error message with generation command hint - Increase DEFAULT_MAX_QUERIES from 128 to 8192 for meaningful throughput tests
67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from scripts.prepare_cam_retrieval_dataset import (
|
|
DEFAULT_MAX_QUERIES,
|
|
dataset_config,
|
|
pack_bits_to_words_le,
|
|
stratified_indices,
|
|
words_le_to_int,
|
|
)
|
|
|
|
|
|
def test_dataset_config_resolves_cifar10_and_cifar100() -> None:
|
|
assert dataset_config("cifar10") == ("uoft-cs/cifar10", "label")
|
|
assert dataset_config("cifar100") == ("uoft-cs/cifar100", "fine_label")
|
|
|
|
|
|
def test_default_query_count_is_large_enough_for_software_throughput() -> None:
|
|
assert DEFAULT_MAX_QUERIES == 8192
|
|
|
|
|
|
def test_dataset_config_rejects_unknown_dataset() -> None:
|
|
try:
|
|
dataset_config("mnist")
|
|
except ValueError as exc:
|
|
assert "dataset must be cifar10 or cifar100" in str(exc)
|
|
else:
|
|
raise AssertionError("expected ValueError")
|
|
|
|
|
|
def test_stratified_indices_balances_labels_and_is_deterministic() -> None:
|
|
labels = [0, 0, 0, 1, 1, 1, 2, 2, 2]
|
|
|
|
first = stratified_indices(labels, total=6, seed=123)
|
|
second = stratified_indices(labels, total=6, seed=123)
|
|
|
|
assert first == second
|
|
selected_labels = [labels[i] for i in first]
|
|
assert selected_labels.count(0) == 2
|
|
assert selected_labels.count(1) == 2
|
|
assert selected_labels.count(2) == 2
|
|
|
|
|
|
def test_stratified_indices_fills_remainder() -> None:
|
|
labels = [0, 0, 0, 1, 1, 1]
|
|
|
|
indices = stratified_indices(labels, total=5, seed=7)
|
|
|
|
assert len(indices) == 5
|
|
assert len(set(indices)) == 5
|
|
|
|
|
|
def test_pack_bits_to_words_le_roundtrip() -> None:
|
|
bits = np.zeros((2, 128), dtype=np.uint8)
|
|
bits[0, 0] = 1
|
|
bits[0, 65] = 1
|
|
bits[1, 63] = 1
|
|
bits[1, 127] = 1
|
|
|
|
words = pack_bits_to_words_le(bits, hash_bits=128)
|
|
|
|
assert words.dtype == np.uint64
|
|
assert words.shape == (2, 2)
|
|
assert words_le_to_int(words[0]) == (1 << 0) | (1 << 65)
|
|
assert words_le_to_int(words[1]) == (1 << 63) | (1 << 127)
|