mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
feat(benchmark): add software CAM retrieval benchmark
Add software-based CAM retrieval benchmark to compare retrieval quality and speed against hardware simulation. Includes experiment documentation with noise sweep analysis on CIFAR-10/100 datasets. - Add sw_retrieval_benchmark.py for software Hamming distance Top-K retrieval - Add test_sw_retrieval_benchmark.py with unit tests for dataset loading and metrics - Add experiment doc (sw_hw_cam_retrieval_benchmark.md) comparing software vs hardware - Document noise sweep impact on retrieval quality at various WRITE_NOISE_RATE values
This commit is contained in:
329
scripts/sw_retrieval_benchmark.py
Normal file
329
scripts/sw_retrieval_benchmark.py
Normal file
@@ -0,0 +1,329 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Callable, Iterable, Sequence
|
||||
|
||||
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]
|
||||
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 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(f"retrieval dataset not found: {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_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,
|
||||
hash_bits=int(rows_words.shape[1] * 64),
|
||||
num_classes=len(set(row_labels)),
|
||||
)
|
||||
|
||||
|
||||
def hamming_similarity_score(query_row: int, stored_row: int, *, width: int) -> int:
|
||||
if width <= 0:
|
||||
raise ValueError("width must be greater than 0")
|
||||
mask = (1 << width) - 1
|
||||
distance = ((int(query_row) ^ int(stored_row)) & mask).bit_count()
|
||||
return int(width - distance)
|
||||
|
||||
|
||||
def match_topk_hamming(query: int, rows: Sequence[int], *, width: int, k: int) -> list[int]:
|
||||
if k <= 0:
|
||||
raise ValueError("k must be greater than 0")
|
||||
scored = [
|
||||
(-hamming_similarity_score(query, row, width=width), row_index)
|
||||
for row_index, row in enumerate(rows)
|
||||
]
|
||||
scored.sort()
|
||||
return [row_index for _, row_index in scored[: min(k, len(scored))]]
|
||||
|
||||
|
||||
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,
|
||||
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")
|
||||
|
||||
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_topk_hamming(query, dataset.rows, width=hash_bits, k=max_k)
|
||||
for query in dataset.queries
|
||||
]
|
||||
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-hamming"
|
||||
|
||||
return {
|
||||
"run_id": resolved_run_id,
|
||||
"mode": "software-hamming",
|
||||
"status": "pass",
|
||||
"params": {
|
||||
"num_rows": len(dataset.rows),
|
||||
"hash_bits": int(hash_bits),
|
||||
"topk_k": max_k,
|
||||
},
|
||||
"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", "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"],
|
||||
"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"- 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` 中的 CAM 行整数执行汉明距离 / 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 software Hamming 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("--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)
|
||||
result = run_benchmark(
|
||||
args.dataset,
|
||||
hash_bits=args.hash_bits,
|
||||
topk_values=topk_values,
|
||||
run_id=args.run_id,
|
||||
)
|
||||
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"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()
|
||||
Reference in New Issue
Block a user