From 1ff9a5f18b3b5de7a741cfd34f235fe712536f84 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Fri, 22 May 2026 21:06:51 +0800 Subject: [PATCH] feat(retrieval-benchmark): add support for external pre-prepared CAM retrieval datasets with recall@k metric - Add just recipes for preparing CIFAR10/100 hash artifacts and running benchmarks - Add CAM_RETRIEVAL_DATASET env var support in Makefile - Add load_retrieval_dataset_npz() to load pre-prepared retrieval datasets - Add label_hits counter and recall@k metric for retrieval evaluation - Rename macro_recall to retrieval_recall to clarify semantics --- .justfile | 16 ++ hw/sim/benchmarks/retrieval/Makefile | 4 + .../retrieval/test_retrieval_benchmark.py | 73 ++++-- scripts/prepare_cam_retrieval_dataset.py | 234 ++++++++++++++++++ tests/test_prepare_cam_retrieval_dataset.py | 61 +++++ 5 files changed, 373 insertions(+), 15 deletions(-) create mode 100644 scripts/prepare_cam_retrieval_dataset.py create mode 100644 tests/test_prepare_cam_retrieval_dataset.py diff --git a/.justfile b/.justfile index 6bba46d..325cb2a 100644 --- a/.justfile +++ b/.justfile @@ -81,3 +81,19 @@ cam-test-retrieval-no-noise: # Run CAM retrieval benchmark with read noise enabled cam-test-retrieval-read-noise: just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 WRITE_NOISE_EN=0 READ_NOISE_EN=1 READ_NOISE_RATE_NUM=1 READ_NOISE_RATE_DEN=100 READ_NOISE_BITS=8" + +# Prepare CIFAR10 hash artifact for CAM retrieval smoke benchmark +cam-prepare-retrieval-cifar10 ROWS="512" QUERIES="128": + just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar10 --num-rows {{ROWS}} --max-queries {{QUERIES}} --compressor-path outputs/hash_compressor.pt" + +# Prepare CIFAR100 hash artifact for CAM retrieval benchmark +cam-prepare-retrieval-cifar100 ROWS="512" QUERIES="128": + just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar100 --num-rows {{ROWS}} --max-queries {{QUERIES}} --compressor-path outputs/hash_compressor.pt" + +# Run CAM retrieval benchmark on a prepared artifact without hardware noise +cam-test-retrieval-artifact DATASET_PATH NUM_ROWS="4096": + just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{NUM_ROWS}} WRITE_NOISE_EN=0 READ_NOISE_EN=0 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}" + +# Run CAM retrieval benchmark on a prepared artifact with read noise enabled +cam-test-retrieval-artifact-read-noise DATASET_PATH NUM_ROWS="4096": + just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{NUM_ROWS}} WRITE_NOISE_EN=0 READ_NOISE_EN=1 READ_NOISE_RATE_NUM=1 READ_NOISE_RATE_DEN=100 READ_NOISE_BITS=8 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}" diff --git a/hw/sim/benchmarks/retrieval/Makefile b/hw/sim/benchmarks/retrieval/Makefile index c300be7..7e77eec 100644 --- a/hw/sim/benchmarks/retrieval/Makefile +++ b/hw/sim/benchmarks/retrieval/Makefile @@ -7,7 +7,11 @@ COCOTB_TEST_MODULES := benchmarks.retrieval.test_retrieval_benchmark VERILOG_SOURCES := $(RTL_CAM_TOP) TOPK_K ?= 5 +NUM_ROWS ?= 4096 WRITE_NOISE_EN ?= 0 READ_NOISE_EN ?= 0 +CAM_RETRIEVAL_DATASET ?= +export CAM_RETRIEVAL_DATASET + include $(SIM_ROOT)/mk/cocotb-common.mk diff --git a/hw/sim/benchmarks/retrieval/test_retrieval_benchmark.py b/hw/sim/benchmarks/retrieval/test_retrieval_benchmark.py index f2db9ac..0cc8b2b 100644 --- a/hw/sim/benchmarks/retrieval/test_retrieval_benchmark.py +++ b/hw/sim/benchmarks/retrieval/test_retrieval_benchmark.py @@ -54,14 +54,16 @@ class MetricAccumulator: 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, exact: bool) -> "MetricAccumulator": + 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, ) @@ -69,15 +71,17 @@ class MetricAccumulator: if self.count == 0: return { "macro_precision": 0.0, - "macro_recall": 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, - "macro_recall": self.recall_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, } @@ -96,6 +100,37 @@ def _flip_exact_bits(rng: np.random.Generator, width: int, n_bits: int) -> int: return mask +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 | os.PathLike[str]) -> RetrievalDataset: + dataset_path = Path(path) + if not dataset_path.is_absolute(): + dataset_path = _project_root() / dataset_path + if not dataset_path.exists(): + raise AssertionError(f"CAM_RETRIEVAL_DATASET not found: {dataset_path}") + loaded = np.load(dataset_path) + rows = [words_le_to_int(words) for words in loaded["rows_words"]] + queries = [words_le_to_int(words) for words in loaded["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, + num_classes=len(set(row_labels)), + positives_per_class=0, + queries_per_class=0, + seed=0, + ) + + def make_clustered_dataset( *, num_rows: int, @@ -195,8 +230,8 @@ def write_outputs(out_dir: Path, result: dict) -> None: "run_id", "mode", "num_rows", "hash_bits", "lanes", "topk_k", "write_noise_en", "read_noise_en", "write_noise_rate_num", "write_noise_rate_den", "read_noise_rate_num", "read_noise_rate_den", - "num_queries", "k", "macro_precision", "macro_recall", "macro_f1", - "exact_match_rate", "status", + "num_queries", "k", "macro_precision", "retrieval_recall", "macro_f1", + "recall@k", "exact_match_rate", "status", ] with metrics_csv.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) @@ -218,8 +253,9 @@ def write_outputs(out_dir: Path, result: dict) -> None: "num_queries": result["dataset"]["num_queries"], "k": int(k), "macro_precision": metrics["macro_precision"], - "macro_recall": metrics["macro_recall"], + "retrieval_recall": metrics["retrieval_recall"], "macro_f1": metrics["macro_f1"], + "recall@k": metrics["recall@k"], "exact_match_rate": metrics["exact_match_rate"], "status": result["status"], } @@ -233,13 +269,13 @@ def write_outputs(out_dir: Path, result: dict) -> None: f"- status: `{result['status']}`", f"- num_queries: `{result['dataset']['num_queries']}`", "", - "| k | macro_precision | macro_recall | macro_f1 | exact_match_rate |", - "|---:|---:|---:|---:|---:|", + "| 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['macro_recall']:.6f} | " - f"{metrics['macro_f1']:.6f} | {metrics['exact_match_rate']:.6f} |" + 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([ "", @@ -270,7 +306,12 @@ async def cam_retrieval_benchmark(dut): if num_rows % lanes != 0: raise AssertionError("Retrieval benchmark requires NUM_ROWS divisible by LANES") - dataset = make_clustered_dataset(num_rows=num_rows, hash_bits=hash_bits) + dataset_path = os.environ.get("CAM_RETRIEVAL_DATASET") + if not dataset_path: + raise AssertionError("CAM_RETRIEVAL_DATASET is required; run scripts/prepare_cam_retrieval_dataset.py first") + dataset = load_retrieval_dataset_npz(dataset_path) + if len(dataset.rows) != num_rows: + raise AssertionError(f"artifact row count {len(dataset.rows)} must equal DUT NUM_ROWS {num_rows}") await write_rows(dut, dataset.rows) accumulators = {k: MetricAccumulator() for k in BENCHMARK_KS} @@ -296,7 +337,9 @@ async def cam_retrieval_benchmark(dut): for k in BENCHMARK_KS: precision, recall, f1 = compute_metrics(dut_topk, dataset.row_labels, query_label, k) exact = dut_topk[:k] == golden_topk[:k] - accumulators[k] = accumulators[k].add(precision, recall, f1, exact) + retrieved_labels = [dataset.row_labels[idx] for idx in dut_topk[:k]] + label_hit = query_label in retrieved_labels + accumulators[k] = accumulators[k].add(precision, recall, f1, label_hit, exact) run_id = os.environ.get("CAM_RETRIEVAL_RUN_ID") or f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-{mode}" result = { @@ -331,9 +374,9 @@ async def cam_retrieval_benchmark(dut): for k in BENCHMARK_KS: metrics = result["metrics"][str(k)] dut._log.info( - "RETRIEVAL_RESULT mode=%s k=%d precision=%.6f recall=%.6f f1=%.6f exact_match=%.6f output_dir=%s", - mode, k, metrics["macro_precision"], metrics["macro_recall"], - metrics["macro_f1"], metrics["exact_match_rate"], + "RETRIEVAL_RESULT mode=%s k=%d precision=%.6f retrieval_recall=%.6f f1=%.6f recall_at_k=%.6f exact_match=%.6f output_dir=%s", + mode, k, metrics["macro_precision"], metrics["retrieval_recall"], + metrics["macro_f1"], metrics["recall@k"], metrics["exact_match_rate"], str(out_dir.relative_to(_project_root())), ) diff --git a/scripts/prepare_cam_retrieval_dataset.py b/scripts/prepare_cam_retrieval_dataset.py new file mode 100644 index 0000000..f86d899 --- /dev/null +++ b/scripts/prepare_cam_retrieval_dataset.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import numpy as np +import torch +import typer +from datasets import load_dataset +from rich.progress import track +from torch.utils.data import DataLoader + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +MINI_NAV_ROOT = PROJECT_ROOT / "mini-nav" +if str(MINI_NAV_ROOT) not in sys.path: + sys.path.insert(0, str(MINI_NAV_ROOT)) + +from compressors.model_loader import get_dino_dim, load_dino_model, load_hash_compressor # noqa: E402 +from utils import get_device # noqa: E402 + +DEFAULT_COMPRESSOR_PATH = Path("outputs/hash_compressor.pt") +DEFAULT_OUTPUT_ROOT = Path("outputs/cam_retrieval_benchmark/datasets") + + +@dataclass(frozen=True) +class PreparedArtifact: + npz_path: Path + metadata_path: Path + num_rows: int + num_queries: int + + +def dataset_config(dataset: str) -> tuple[str, str]: + if dataset == "cifar10": + return "uoft-cs/cifar10", "label" + if dataset == "cifar100": + return "uoft-cs/cifar100", "fine_label" + raise ValueError("dataset must be cifar10 or cifar100") + + +def stratified_indices(labels: list[int], *, total: int, seed: int) -> list[int]: + if total <= 0: + raise ValueError("total must be greater than 0") + if total > len(labels): + raise ValueError("total cannot exceed number of labels") + + rng = np.random.default_rng(seed) + by_label: dict[int, list[int]] = {} + for idx, label in enumerate(labels): + by_label.setdefault(int(label), []).append(idx) + + label_ids = sorted(by_label) + base = total // len(label_ids) + remainder = total % len(label_ids) + selected: list[int] = [] + + for offset, label in enumerate(label_ids): + want = base + (1 if offset < remainder else 0) + candidates = list(by_label[label]) + if want > len(candidates): + raise ValueError(f"not enough examples for label {label}: need {want}, have {len(candidates)}") + chosen = rng.choice(candidates, size=want, replace=False).tolist() + selected.extend(int(i) for i in chosen) + + selected.sort() + return selected + + +def pack_bits_to_words_le(bits: np.ndarray, *, hash_bits: int) -> np.ndarray: + if bits.ndim != 2: + raise ValueError("bits must have shape [N, hash_bits]") + if bits.shape[1] != hash_bits: + raise ValueError(f"expected {hash_bits} bits, got {bits.shape[1]}") + if hash_bits % 64 != 0: + raise ValueError("hash_bits must be divisible by 64") + + words = np.zeros((bits.shape[0], hash_bits // 64), dtype=np.uint64) + bits_u8 = bits.astype(np.uint8, copy=False) + for word_idx in range(hash_bits // 64): + word = np.zeros(bits.shape[0], dtype=np.uint64) + for bit in range(64): + word |= bits_u8[:, word_idx * 64 + bit].astype(np.uint64) << np.uint64(bit) + words[:, word_idx] = word + return words + + +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 + + +@torch.no_grad() +def encode_dataset_bits( + dataset, + *, + img_column: str, + batch_size: int, + dino_model: str, + compressor_path: Path, + hash_bits: int, +) -> np.ndarray: + if not compressor_path.exists(): + raise FileNotFoundError(f"missing HashCompressor weights: {compressor_path}") + + device = get_device() + processor, dino = load_dino_model(dino_model) + compressor = load_hash_compressor( + input_dim=get_dino_dim(dino_model), + hash_bits=hash_bits, + compressor_path=str(compressor_path), + ).to(device) + compressor.eval() + + loader = DataLoader(dataset.with_format("torch"), batch_size=batch_size, shuffle=False, num_workers=4) + all_bits: list[np.ndarray] = [] + for batch in track(loader, description="Encoding CIFAR hash bits"): + images = batch[img_column] + inputs = processor(images=images, return_tensors="pt").to(device) + tokens = dino(**inputs).last_hidden_state + bits = compressor.encode(tokens).detach().cpu().numpy().astype(np.uint8) + all_bits.append(bits) + return np.concatenate(all_bits, axis=0) + + +def prepare_artifact( + *, + dataset_name: Literal["cifar10", "cifar100"], + num_rows: int, + max_queries: int, + compressor_path: Path, + output_root: Path, + dino_model: str = "facebook/dinov2-large", + hash_bits: int = 512, + batch_size: int = 64, + seed: int = 20260522, +) -> PreparedArtifact: + hf_id, label_column = dataset_config(dataset_name) + loaded = load_dataset(hf_id) + train = loaded["train"] + test = loaded["test"] if "test" in loaded else loaded["validation"] + + row_indices = stratified_indices([int(x) for x in train[label_column]], total=num_rows, seed=seed) + query_indices = stratified_indices([int(x) for x in test[label_column]], total=max_queries, seed=seed + 1) + + row_subset = train.select(row_indices) + query_subset = test.select(query_indices) + + row_bits = encode_dataset_bits( + row_subset, + img_column="img", + batch_size=batch_size, + dino_model=dino_model, + compressor_path=compressor_path, + hash_bits=hash_bits, + ) + query_bits = encode_dataset_bits( + query_subset, + img_column="img", + batch_size=batch_size, + dino_model=dino_model, + compressor_path=compressor_path, + hash_bits=hash_bits, + ) + + rows_words = pack_bits_to_words_le(row_bits, hash_bits=hash_bits) + queries_words = pack_bits_to_words_le(query_bits, hash_bits=hash_bits) + row_labels = np.array([int(x) for x in row_subset[label_column]], dtype=np.int64) + query_labels = np.array([int(x) for x in query_subset[label_column]], dtype=np.int64) + + output_root.mkdir(parents=True, exist_ok=True) + stem = f"{dataset_name}_hash{hash_bits}_rows{num_rows}_queries{max_queries}" + npz_path = output_root / f"{stem}.npz" + metadata_path = output_root / f"{stem}.json" + + np.savez_compressed( + npz_path, + rows_words=rows_words, + row_labels=row_labels, + queries_words=queries_words, + query_labels=query_labels, + ) + metadata = { + "dataset": dataset_name, + "hf_id": hf_id, + "label_column": label_column, + "compressor_path": str(compressor_path), + "dino_model": dino_model, + "hash_bits": hash_bits, + "bit_source": "HashCompressor.encode(tokens)", + "bit_values": "{0,1} int32", + "num_rows": int(num_rows), + "num_queries": int(max_queries), + "rows_source_split": "train", + "queries_source_split": "test", + "seed": int(seed), + } + metadata_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return PreparedArtifact(npz_path=npz_path, metadata_path=metadata_path, num_rows=num_rows, num_queries=max_queries) + + +def main( + dataset: Literal["cifar10", "cifar100"] = typer.Option("cifar100"), + num_rows: int = typer.Option(512, min=5), + max_queries: int = typer.Option(128, min=1), + compressor_path: Path = typer.Option(DEFAULT_COMPRESSOR_PATH), + output_root: Path = typer.Option(DEFAULT_OUTPUT_ROOT), + dino_model: str = typer.Option("facebook/dinov2-large"), + hash_bits: int = typer.Option(512), + batch_size: int = typer.Option(64, min=1), + seed: int = typer.Option(20260522), +) -> None: + artifact = prepare_artifact( + dataset_name=dataset, + num_rows=num_rows, + max_queries=max_queries, + compressor_path=compressor_path, + output_root=output_root, + dino_model=dino_model, + hash_bits=hash_bits, + batch_size=batch_size, + seed=seed, + ) + print(f"CAM_RETRIEVAL_DATASET={artifact.npz_path}") + print(f"CAM_RETRIEVAL_METADATA={artifact.metadata_path}") + + +if __name__ == "__main__": + typer.run(main) diff --git a/tests/test_prepare_cam_retrieval_dataset.py b/tests/test_prepare_cam_retrieval_dataset.py new file mode 100644 index 0000000..c26dcaa --- /dev/null +++ b/tests/test_prepare_cam_retrieval_dataset.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import numpy as np + +from scripts.prepare_cam_retrieval_dataset import ( + 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_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)