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
437 lines
16 KiB
Python
437 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import re
|
|
import sys
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Callable, Iterable
|
|
|
|
import numpy as np
|
|
|
|
|
|
BENCHMARK_KS = (1, 5)
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
HW_SIM_DIR = PROJECT_ROOT / "hw" / "sim"
|
|
if str(HW_SIM_DIR) not in sys.path:
|
|
sys.path.insert(0, str(HW_SIM_DIR))
|
|
|
|
from model.ref_model import match_topk as ref_match_topk # noqa: E402
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetrievalDataset:
|
|
rows: list[int]
|
|
row_labels: list[int]
|
|
queries: list[int]
|
|
query_labels: list[int]
|
|
rows_words: np.ndarray
|
|
queries_words: np.ndarray
|
|
hash_bits: int
|
|
num_classes: int
|
|
positives_per_class: int = 0
|
|
queries_per_class: int = 0
|
|
seed: int = 0
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MetricAccumulator:
|
|
precision_sum: float = 0.0
|
|
recall_sum: float = 0.0
|
|
f1_sum: float = 0.0
|
|
exact_matches: int = 0
|
|
label_hits: int = 0
|
|
count: int = 0
|
|
|
|
def add(self, precision: float, recall: float, f1: float, label_hit: bool, exact: bool) -> "MetricAccumulator":
|
|
return MetricAccumulator(
|
|
precision_sum=self.precision_sum + precision,
|
|
recall_sum=self.recall_sum + recall,
|
|
f1_sum=self.f1_sum + f1,
|
|
exact_matches=self.exact_matches + int(exact),
|
|
label_hits=self.label_hits + int(label_hit),
|
|
count=self.count + 1,
|
|
)
|
|
|
|
def as_dict(self) -> dict[str, float]:
|
|
if self.count == 0:
|
|
return {
|
|
"macro_precision": 0.0,
|
|
"retrieval_recall": 0.0,
|
|
"macro_f1": 0.0,
|
|
"exact_match_rate": 0.0,
|
|
"recall@k": 0.0,
|
|
}
|
|
return {
|
|
"macro_precision": self.precision_sum / self.count,
|
|
"retrieval_recall": self.recall_sum / self.count,
|
|
"macro_f1": self.f1_sum / self.count,
|
|
"exact_match_rate": self.exact_matches / self.count,
|
|
"recall@k": self.label_hits / self.count,
|
|
}
|
|
|
|
|
|
def project_root() -> Path:
|
|
return PROJECT_ROOT
|
|
|
|
|
|
def missing_dataset_message(dataset_path: Path) -> str:
|
|
message = f"retrieval dataset not found: {dataset_path}"
|
|
match = re.fullmatch(
|
|
r"(?P<dataset>cifar10|cifar100)_hash(?P<hash_bits>\d+)_rows(?P<num_rows>\d+)_queries(?P<max_queries>\d+)\.npz",
|
|
dataset_path.name,
|
|
)
|
|
if match is None:
|
|
return message
|
|
|
|
groups = match.groupdict()
|
|
command = (
|
|
"python scripts/prepare_cam_retrieval_dataset.py "
|
|
f"--dataset {groups['dataset']} "
|
|
f"--num-rows {groups['num_rows']} "
|
|
f"--max-queries {groups['max_queries']} "
|
|
f"--hash-bits {groups['hash_bits']}"
|
|
)
|
|
return (
|
|
f"{message}\n"
|
|
"The requested benchmark artifact has not been generated yet. "
|
|
"Create it first, then rerun this benchmark:\n"
|
|
f" {command}"
|
|
)
|
|
|
|
|
|
def words_le_to_int(words: np.ndarray) -> int:
|
|
value = 0
|
|
for idx, word in enumerate(words.tolist()):
|
|
value |= int(word) << (64 * idx)
|
|
return value
|
|
|
|
|
|
def load_retrieval_dataset_npz(path: str | Path) -> RetrievalDataset:
|
|
dataset_path = Path(path)
|
|
if not dataset_path.is_absolute():
|
|
dataset_path = project_root() / dataset_path
|
|
if not dataset_path.exists():
|
|
raise FileNotFoundError(missing_dataset_message(dataset_path))
|
|
|
|
loaded = np.load(dataset_path)
|
|
rows_words = loaded["rows_words"]
|
|
queries_words = loaded["queries_words"]
|
|
if rows_words.ndim != 2 or queries_words.ndim != 2:
|
|
raise ValueError("rows_words and queries_words must have shape [N, words]")
|
|
if rows_words.shape[1] != queries_words.shape[1]:
|
|
raise ValueError("rows_words and queries_words must use the same word width")
|
|
|
|
rows_words = np.asarray(rows_words, dtype=np.uint64)
|
|
queries_words = np.asarray(queries_words, dtype=np.uint64)
|
|
rows = [words_le_to_int(words) for words in rows_words]
|
|
queries = [words_le_to_int(words) for words in queries_words]
|
|
row_labels = [int(x) for x in loaded["row_labels"].tolist()]
|
|
query_labels = [int(x) for x in loaded["query_labels"].tolist()]
|
|
return RetrievalDataset(
|
|
rows=rows,
|
|
row_labels=row_labels,
|
|
queries=queries,
|
|
query_labels=query_labels,
|
|
rows_words=rows_words,
|
|
queries_words=queries_words,
|
|
hash_bits=int(rows_words.shape[1] * 64),
|
|
num_classes=len(set(row_labels)),
|
|
)
|
|
|
|
|
|
def match_topk_numpy(query_words: np.ndarray, rows_words: np.ndarray, *, width: int, k: int) -> list[int]:
|
|
if k <= 0:
|
|
raise ValueError("k must be greater than 0")
|
|
if width <= 0:
|
|
raise ValueError("width must be greater than 0")
|
|
query_words = np.asarray(query_words, dtype=np.uint64)
|
|
rows_words = np.asarray(rows_words, dtype=np.uint64)
|
|
if rows_words.ndim != 2:
|
|
raise ValueError("rows_words must have shape [N, words]")
|
|
if query_words.shape != (rows_words.shape[1],):
|
|
raise ValueError("query_words must have shape [words]")
|
|
|
|
return _match_topk_numpy_batch(query_words[np.newaxis, :], rows_words, width=width, k=k)[0]
|
|
|
|
|
|
def _match_topk_numpy_batch(
|
|
queries_words: np.ndarray,
|
|
rows_words: np.ndarray,
|
|
*,
|
|
width: int,
|
|
k: int,
|
|
) -> list[list[int]]:
|
|
if k <= 0:
|
|
raise ValueError("k must be greater than 0")
|
|
if width <= 0:
|
|
raise ValueError("width must be greater than 0")
|
|
queries_words = np.asarray(queries_words, dtype=np.uint64)
|
|
rows_words = np.asarray(rows_words, dtype=np.uint64)
|
|
if queries_words.ndim != 2:
|
|
raise ValueError("queries_words must have shape [Q, words]")
|
|
if rows_words.ndim != 2:
|
|
raise ValueError("rows_words must have shape [N, words]")
|
|
if queries_words.shape[1] != rows_words.shape[1]:
|
|
raise ValueError("queries_words and rows_words must use the same word width")
|
|
|
|
xor = np.bitwise_xor(queries_words[:, np.newaxis, :], rows_words[np.newaxis, :, :])
|
|
distances = np.bitwise_count(xor).sum(axis=2, dtype=np.int64)
|
|
scores = int(width) - distances
|
|
row_indices = np.arange(rows_words.shape[0], dtype=np.int64)
|
|
topk_count = min(k, rows_words.shape[0])
|
|
return [
|
|
[int(idx) for idx in np.lexsort((row_indices, -score_row))[:topk_count]]
|
|
for score_row in scores
|
|
]
|
|
|
|
|
|
def match_all_topk_numpy(
|
|
queries_words: np.ndarray,
|
|
rows_words: np.ndarray,
|
|
*,
|
|
width: int,
|
|
k: int,
|
|
workers: int = 1,
|
|
) -> list[list[int]]:
|
|
if workers <= 0:
|
|
raise ValueError("workers must be greater than 0")
|
|
queries_words = np.asarray(queries_words, dtype=np.uint64)
|
|
if workers == 1 or len(queries_words) <= 1:
|
|
return _match_topk_numpy_batch(queries_words, rows_words, width=width, k=k)
|
|
|
|
chunks = [
|
|
chunk
|
|
for chunk in np.array_split(queries_words, min(workers, len(queries_words)))
|
|
if len(chunk)
|
|
]
|
|
with ThreadPoolExecutor(max_workers=workers) as executor:
|
|
futures = [
|
|
executor.submit(_match_topk_numpy_batch, chunk, rows_words, width=width, k=k)
|
|
for chunk in chunks
|
|
]
|
|
parts = [future.result() for future in futures]
|
|
return [topk for part in parts for topk in part]
|
|
|
|
|
|
def compute_metrics(topk_indices: list[int], row_labels: list[int], query_label: int, k: int) -> tuple[float, float, float]:
|
|
retrieved = topk_indices[:k]
|
|
relevant = {idx for idx, label in enumerate(row_labels) if label == query_label}
|
|
tp = len(set(retrieved) & relevant)
|
|
precision = tp / float(k)
|
|
recall = tp / float(len(relevant)) if relevant else 0.0
|
|
f1 = 0.0 if precision + recall == 0 else (2.0 * precision * recall) / (precision + recall)
|
|
return precision, recall, f1
|
|
|
|
|
|
def _normalized_topk_values(topk_values: Iterable[int]) -> tuple[int, ...]:
|
|
values = tuple(sorted({int(k) for k in topk_values}))
|
|
if not values or values[0] <= 0:
|
|
raise ValueError("topk_values must contain positive integers")
|
|
return values
|
|
|
|
|
|
def run_benchmark(
|
|
dataset_path: str | Path,
|
|
*,
|
|
hash_bits: int = 512,
|
|
topk_values: Iterable[int] = BENCHMARK_KS,
|
|
run_id: str | None = None,
|
|
workers: int = 1,
|
|
timer_ns: Callable[[], int] = time.perf_counter_ns,
|
|
) -> dict:
|
|
dataset = load_retrieval_dataset_npz(dataset_path)
|
|
if not dataset.rows:
|
|
raise ValueError("cannot benchmark an empty row set")
|
|
if hash_bits != dataset.hash_bits:
|
|
raise ValueError(
|
|
f"hash_bits={hash_bits} does not match dataset width {dataset.hash_bits}"
|
|
)
|
|
if len(dataset.queries) != len(dataset.query_labels):
|
|
raise ValueError("queries and query_labels must have the same length")
|
|
if workers <= 0:
|
|
raise ValueError("workers must be greater than 0")
|
|
|
|
ks = _normalized_topk_values(topk_values)
|
|
max_k = max(ks)
|
|
if max_k > len(dataset.rows):
|
|
raise ValueError("topk_values cannot exceed the number of dataset rows")
|
|
|
|
start_ns = timer_ns()
|
|
all_topk = match_all_topk_numpy(
|
|
dataset.queries_words,
|
|
dataset.rows_words,
|
|
width=hash_bits,
|
|
k=max_k,
|
|
workers=workers,
|
|
)
|
|
end_ns = timer_ns()
|
|
|
|
golden_topk = [
|
|
ref_match_topk(query, dataset.rows, width=hash_bits, k=max_k)[0]
|
|
for query in dataset.queries
|
|
]
|
|
|
|
accumulators = {k: MetricAccumulator() for k in ks}
|
|
for topk_indices, golden_indices, query_label in zip(all_topk, golden_topk, dataset.query_labels):
|
|
for k in ks:
|
|
precision, recall, f1 = compute_metrics(topk_indices, dataset.row_labels, query_label, k)
|
|
retrieved_labels = [dataset.row_labels[idx] for idx in topk_indices[:k]]
|
|
label_hit = query_label in retrieved_labels
|
|
exact = topk_indices[:k] == golden_indices[:k]
|
|
accumulators[k] = accumulators[k].add(precision, recall, f1, label_hit, exact)
|
|
|
|
metrics = {str(k): accumulators[k].as_dict() for k in ks}
|
|
|
|
elapsed_ns = max(0, int(end_ns - start_ns))
|
|
num_queries = len(dataset.queries)
|
|
ns_per_query = (elapsed_ns / float(num_queries)) if num_queries else 0.0
|
|
qps = (1_000_000_000.0 / ns_per_query) if ns_per_query > 0 else 0.0
|
|
resolved_run_id = run_id or f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-software-numpy"
|
|
|
|
return {
|
|
"run_id": resolved_run_id,
|
|
"mode": "software-numpy",
|
|
"status": "pass",
|
|
"params": {
|
|
"num_rows": len(dataset.rows),
|
|
"hash_bits": int(hash_bits),
|
|
"topk_k": max_k,
|
|
"workers": int(workers),
|
|
"engine": "numpy",
|
|
},
|
|
"dataset": {
|
|
"num_classes": dataset.num_classes,
|
|
"positives_per_class": dataset.positives_per_class,
|
|
"queries_per_class": dataset.queries_per_class,
|
|
"num_queries": num_queries,
|
|
"seed": dataset.seed,
|
|
},
|
|
"metrics": metrics,
|
|
"performance": {
|
|
"total_elapsed_ns": elapsed_ns,
|
|
"total_elapsed_sec": elapsed_ns / 1_000_000_000.0,
|
|
"ns_per_query": ns_per_query,
|
|
"queries_per_second": qps,
|
|
},
|
|
}
|
|
|
|
|
|
def write_outputs(out_dir: Path, result: dict) -> None:
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
(out_dir / "metrics.json").write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
|
|
|
fieldnames = [
|
|
"run_id", "mode", "num_rows", "hash_bits", "topk_k", "workers", "engine", "num_queries", "k",
|
|
"macro_precision", "retrieval_recall", "macro_f1", "recall@k", "exact_match_rate",
|
|
"total_elapsed_ns", "ns_per_query", "queries_per_second", "status",
|
|
]
|
|
with (out_dir / "metrics.csv").open("w", newline="", encoding="utf-8") as f:
|
|
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
for k, metrics in result["metrics"].items():
|
|
writer.writerow({
|
|
"run_id": result["run_id"],
|
|
"mode": result["mode"],
|
|
"num_rows": result["params"]["num_rows"],
|
|
"hash_bits": result["params"]["hash_bits"],
|
|
"topk_k": result["params"]["topk_k"],
|
|
"workers": result["params"].get("workers", 1),
|
|
"engine": result["params"].get("engine", "numpy"),
|
|
"num_queries": result["dataset"]["num_queries"],
|
|
"k": int(k),
|
|
"macro_precision": metrics["macro_precision"],
|
|
"retrieval_recall": metrics["retrieval_recall"],
|
|
"macro_f1": metrics["macro_f1"],
|
|
"recall@k": metrics["recall@k"],
|
|
"exact_match_rate": metrics["exact_match_rate"],
|
|
"total_elapsed_ns": result["performance"]["total_elapsed_ns"],
|
|
"ns_per_query": result["performance"]["ns_per_query"],
|
|
"queries_per_second": result["performance"]["queries_per_second"],
|
|
"status": result["status"],
|
|
})
|
|
|
|
lines = [
|
|
"# Software CAM Retrieval Benchmark Summary",
|
|
"",
|
|
f"- run_id: `{result['run_id']}`",
|
|
f"- mode: `{result['mode']}`",
|
|
f"- status: `{result['status']}`",
|
|
f"- num_queries: `{result['dataset']['num_queries']}`",
|
|
f"- workers: `{result['params'].get('workers', 1)}`",
|
|
f"- engine: `{result['params'].get('engine', 'numpy')}`",
|
|
f"- ns_per_query: `{result['performance']['ns_per_query']}`",
|
|
f"- queries_per_second: `{result['performance']['queries_per_second']}`",
|
|
"",
|
|
"| k | macro_precision | retrieval_recall | macro_f1 | recall@k | exact_match_rate |",
|
|
"|---:|---:|---:|---:|---:|---:|",
|
|
]
|
|
for k, metrics in result["metrics"].items():
|
|
lines.append(
|
|
f"| {k} | {metrics['macro_precision']:.6f} | {metrics['retrieval_recall']:.6f} | "
|
|
f"{metrics['macro_f1']:.6f} | {metrics['recall@k']:.6f} | {metrics['exact_match_rate']:.6f} |"
|
|
)
|
|
lines.extend([
|
|
"",
|
|
"说明：软件路径直接对 `.npz` 中的 little-endian uint64 words 使用 NumPy bitwise_count 执行汉明距离 / XNOR-popcount 检索，不使用软件 CAM 时序仿真。",
|
|
])
|
|
(out_dir / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
def output_dir_for(run_id: str, output_root: Path) -> Path:
|
|
return output_root / run_id
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Run NumPy software CAM retrieval benchmark.")
|
|
parser.add_argument("--dataset", required=True, help="Prepared CAM retrieval .npz artifact")
|
|
parser.add_argument("--hash-bits", type=int, default=512, help="Hash width in bits")
|
|
parser.add_argument("--topk-k", type=int, default=5, help="Maximum Top-K to report; reports k=1 and this value")
|
|
parser.add_argument("--workers", type=int, default=1, help="Number of software query worker threads")
|
|
parser.add_argument("--run-id", default=None, help="Output run id")
|
|
parser.add_argument(
|
|
"--output-root",
|
|
type=Path,
|
|
default=Path("outputs/sw_retrieval_benchmark"),
|
|
help="Directory under which the run directory is written",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
topk_values = (1,) if args.topk_k == 1 else (1, args.topk_k)
|
|
try:
|
|
result = run_benchmark(
|
|
args.dataset,
|
|
hash_bits=args.hash_bits,
|
|
topk_values=topk_values,
|
|
run_id=args.run_id,
|
|
workers=args.workers,
|
|
)
|
|
except FileNotFoundError as exc:
|
|
print(str(exc), file=sys.stderr)
|
|
raise SystemExit(2) from None
|
|
out_dir = output_dir_for(result["run_id"], args.output_root)
|
|
write_outputs(out_dir, result)
|
|
print(
|
|
"SW_RETRIEVAL_RESULT "
|
|
f"run_id={result['run_id']} "
|
|
f"num_rows={result['params']['num_rows']} "
|
|
f"hash_bits={result['params']['hash_bits']} "
|
|
f"workers={result['params']['workers']} "
|
|
f"num_queries={result['dataset']['num_queries']} "
|
|
f"ns_per_query={result['performance']['ns_per_query']:.3f} "
|
|
f"queries_per_second={result['performance']['queries_per_second']:.3f} "
|
|
f"output_dir={out_dir}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|