feat: 实现明文版hnsw算法和最优参数
This commit is contained in:
338
test_fhe_hnsw.py
Executable file
338
test_fhe_hnsw.py
Executable file
@@ -0,0 +1,338 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
FHE HNSW并行测试脚本 - 多进程版本
|
||||
测试密文态HNSW实现的稳定性和准确率
|
||||
支持多进程并行执行，充分利用服务器CPU资源
|
||||
运行100次，记录≥90%准确率的成功率和结果长度分布
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import json
|
||||
import time
|
||||
import os
|
||||
import multiprocessing as mp
|
||||
from pathlib import Path
|
||||
from collections import defaultdict
|
||||
|
||||
# 正确答案
|
||||
CORRECT_ANSWER = [93, 94, 90, 27, 87, 50, 47, 40, 78, 28]
|
||||
|
||||
|
||||
def calculate_accuracy(result, correct):
|
||||
"""计算准确率：匹配元素数量 / 总元素数量"""
|
||||
matches = len(set(result) & set(correct))
|
||||
return matches / len(correct) * 100
|
||||
|
||||
|
||||
def run_single_test(test_id, data_bit_width="12", timeout_minutes=30):
|
||||
"""运行单次FHE HNSW测试（供多进程调用）"""
|
||||
# 为每个进程创建独立的输出文件
|
||||
output_file = f"./test_fhe_output_{test_id}_{os.getpid()}.jsonl"
|
||||
|
||||
cmd = [
|
||||
"./enc",
|
||||
"--algorithm",
|
||||
"hnsw",
|
||||
"--data-bit-width",
|
||||
data_bit_width,
|
||||
"--predictions",
|
||||
output_file,
|
||||
]
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd, capture_output=True, text=True, timeout=timeout_minutes * 60
|
||||
)
|
||||
|
||||
test_time = time.time() - start_time
|
||||
|
||||
if result.returncode != 0:
|
||||
return {
|
||||
"test_id": test_id,
|
||||
"result": None,
|
||||
"error": f"Command failed: {result.stderr[:200]}",
|
||||
"time_minutes": test_time / 60,
|
||||
}
|
||||
|
||||
# 读取结果
|
||||
if not Path(output_file).exists():
|
||||
return {
|
||||
"test_id": test_id,
|
||||
"result": None,
|
||||
"error": "Output file not found",
|
||||
"time_minutes": test_time / 60,
|
||||
}
|
||||
|
||||
with open(output_file, "r") as f:
|
||||
output = json.load(f)
|
||||
answer = output["answer"]
|
||||
|
||||
# 清理临时文件
|
||||
Path(output_file).unlink(missing_ok=True)
|
||||
|
||||
# 计算准确率
|
||||
accuracy = calculate_accuracy(answer, CORRECT_ANSWER)
|
||||
|
||||
return {
|
||||
"test_id": test_id,
|
||||
"result": answer,
|
||||
"length": len(answer),
|
||||
"accuracy": accuracy,
|
||||
"success": accuracy >= 90,
|
||||
"time_minutes": test_time / 60,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
# 清理临时文件
|
||||
Path(output_file).unlink(missing_ok=True)
|
||||
return {
|
||||
"test_id": test_id,
|
||||
"result": None,
|
||||
"error": "timeout",
|
||||
"time_minutes": timeout_minutes,
|
||||
}
|
||||
except Exception as e:
|
||||
# 清理临时文件
|
||||
Path(output_file).unlink(missing_ok=True)
|
||||
return {
|
||||
"test_id": test_id,
|
||||
"result": None,
|
||||
"error": str(e),
|
||||
"time_minutes": (time.time() - start_time) / 60,
|
||||
}
|
||||
|
||||
|
||||
def print_progress(completed, total, start_time):
|
||||
"""打印进度信息"""
|
||||
if completed == 0:
|
||||
return
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
avg_time = elapsed / completed
|
||||
remaining_time = avg_time * (total - completed)
|
||||
|
||||
print(
|
||||
f"\r🔄 进度: {completed}/{total} ({completed/total*100:.1f}%) "
|
||||
f"| 已用时: {elapsed/3600:.1f}h | 预计剩余: {remaining_time/3600:.1f}h",
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""主测试函数"""
|
||||
print("🚀 FHE HNSW 并行批量测试脚本")
|
||||
print(f"🎯 正确答案: {CORRECT_ANSWER}")
|
||||
print("📊 运行100次测试，记录准确率和结果长度分布")
|
||||
print("🔧 支持多进程并行执行")
|
||||
|
||||
# 检查enc二进制文件是否存在
|
||||
if not Path("./enc").exists():
|
||||
print("\n❌ 找不到 ./enc 二进制文件，请先编译项目:")
|
||||
print(" cargo build --release --bin enc")
|
||||
return
|
||||
|
||||
# 获取CPU核心数并设置进程数
|
||||
cpu_cores = mp.cpu_count()
|
||||
# 考虑到FHE运算的内存密集性，使用核心数的一半避免内存不足
|
||||
num_processes = 8
|
||||
timeout_minutes = 45 # 增加超时时间到45分钟
|
||||
|
||||
print(f"💻 使用 {num_processes} 个并行进程")
|
||||
print(f"⏰ 单次测试超时时间: {timeout_minutes} 分钟")
|
||||
print(
|
||||
f"⏱️ 预计总时间: {100 * 15 / num_processes / 60:.1f}-{100 * 25 / num_processes / 60:.1f} 小时"
|
||||
)
|
||||
print()
|
||||
|
||||
print("=" * 80)
|
||||
print("🔬 开始FHE HNSW并行测试 (data-bit-width=12)")
|
||||
print("=" * 80)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# 创建进程池并执行测试
|
||||
test_ids = list(range(1, 101)) # 1-100
|
||||
|
||||
with mp.Pool(processes=num_processes) as pool:
|
||||
# 创建异步任务
|
||||
async_results = []
|
||||
for test_id in test_ids:
|
||||
async_result = pool.apply_async(
|
||||
run_single_test, (test_id, "12", timeout_minutes)
|
||||
)
|
||||
async_results.append(async_result)
|
||||
|
||||
# 收集结果并显示进度
|
||||
results = []
|
||||
completed = 0
|
||||
|
||||
print_progress(completed, len(test_ids), start_time)
|
||||
|
||||
for async_result in async_results:
|
||||
try:
|
||||
result = async_result.get(
|
||||
timeout=timeout_minutes * 60 + 60
|
||||
) # 额外1分钟缓冲
|
||||
results.append(result)
|
||||
completed += 1
|
||||
print_progress(completed, len(test_ids), start_time)
|
||||
except mp.TimeoutError:
|
||||
# 进程级别超时
|
||||
results.append(
|
||||
{
|
||||
"test_id": len(results) + 1,
|
||||
"result": None,
|
||||
"error": "process_timeout",
|
||||
"time_minutes": timeout_minutes,
|
||||
}
|
||||
)
|
||||
completed += 1
|
||||
print_progress(completed, len(test_ids), start_time)
|
||||
|
||||
print() # 换行
|
||||
total_elapsed = time.time() - start_time
|
||||
|
||||
# 分析结果
|
||||
print("\n" + "=" * 80)
|
||||
print("📈 测试结果分析")
|
||||
print("=" * 80)
|
||||
|
||||
# 统计变量
|
||||
valid_results = [r for r in results if r["result"] is not None]
|
||||
success_count = sum(1 for r in valid_results if r["success"])
|
||||
length_distribution = defaultdict(int)
|
||||
error_distribution = defaultdict(int)
|
||||
|
||||
# 统计错误类型
|
||||
for r in results:
|
||||
if r["error"]:
|
||||
error_type = r["error"]
|
||||
if "timeout" in error_type.lower():
|
||||
error_distribution["timeout"] += 1
|
||||
elif "failed" in error_type.lower():
|
||||
error_distribution["failed"] += 1
|
||||
else:
|
||||
error_distribution["other"] += 1
|
||||
else:
|
||||
length_distribution[r["length"]] += 1
|
||||
|
||||
# 基本统计
|
||||
total_tests = len(results)
|
||||
valid_tests = len(valid_results)
|
||||
|
||||
if valid_tests == 0:
|
||||
print("❌ 没有有效的测试结果")
|
||||
return
|
||||
|
||||
success_rate = success_count / valid_tests * 100
|
||||
avg_accuracy = sum(r["accuracy"] for r in valid_results) / valid_tests
|
||||
avg_length = sum(r["length"] for r in valid_results) / valid_tests
|
||||
avg_time_per_test = sum(r["time_minutes"] for r in results) / total_tests
|
||||
|
||||
print(f"总测试次数: {total_tests}")
|
||||
print(f"有效测试次数: {valid_tests}")
|
||||
print(f"成功次数 (≥90%准确率): {success_count}")
|
||||
print(f"成功率: {success_rate:.1f}%")
|
||||
print(f"平均准确率: {avg_accuracy:.1f}%")
|
||||
print(f"平均结果长度: {avg_length:.1f}")
|
||||
print(f"平均每次测试时间: {avg_time_per_test:.1f}分钟")
|
||||
print(f"总测试时间: {total_elapsed/3600:.1f}小时")
|
||||
print(f"并行加速比: {100 * avg_time_per_test / 60 / (total_elapsed/3600):.1f}x")
|
||||
|
||||
# 错误统计
|
||||
if error_distribution:
|
||||
print("\n❌ 错误分布:")
|
||||
for error_type, count in error_distribution.items():
|
||||
print(f" {error_type}: {count}次")
|
||||
|
||||
# 结果长度分布
|
||||
if length_distribution:
|
||||
print("\n📊 结果长度分布:")
|
||||
for length in sorted(length_distribution.keys()):
|
||||
count = length_distribution[length]
|
||||
percentage = count / valid_tests * 100
|
||||
bar = "█" * (count // 2) if count > 0 else ""
|
||||
print(f" 长度 {length:2d}: {count:3d}次 ({percentage:5.1f}%) {bar}")
|
||||
|
||||
# 结论
|
||||
print()
|
||||
if success_rate >= 50:
|
||||
print("✅ 测试通过! FHE HNSW实现稳定性良好")
|
||||
print(f" - 成功率: {success_rate:.1f}% (≥50%)")
|
||||
print(f" - 平均准确率: {avg_accuracy:.1f}%")
|
||||
if avg_length >= 9.5:
|
||||
print(f" - 平均结果长度: {avg_length:.1f} (接近10)")
|
||||
else:
|
||||
print(f" - 平均结果长度: {avg_length:.1f} (需要改进)")
|
||||
else:
|
||||
print("❌ 测试未通过，需要进一步优化")
|
||||
print(f" - 成功率: {success_rate:.1f}% (<50%)")
|
||||
print(f" - 平均准确率: {avg_accuracy:.1f}%")
|
||||
|
||||
# 保存详细结果
|
||||
report_file = "fhe_hnsw_parallel_test_report.json"
|
||||
with open(report_file, "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"summary": {
|
||||
"total_tests": total_tests,
|
||||
"valid_tests": valid_tests,
|
||||
"success_count": success_count,
|
||||
"success_rate": success_rate,
|
||||
"avg_accuracy": avg_accuracy,
|
||||
"avg_length": avg_length,
|
||||
"avg_time_minutes": avg_time_per_test,
|
||||
"total_time_hours": total_elapsed / 3600,
|
||||
"num_processes": num_processes,
|
||||
"cpu_cores": cpu_cores,
|
||||
"speedup": 100 * avg_time_per_test / 60 / (total_elapsed / 3600),
|
||||
},
|
||||
"length_distribution": dict(length_distribution),
|
||||
"error_distribution": dict(error_distribution),
|
||||
"detailed_results": results,
|
||||
"correct_answer": CORRECT_ANSWER,
|
||||
},
|
||||
f,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
print(f"\n📁 详细报告已保存到: {report_file}")
|
||||
|
||||
# 显示最好和最坏的几个结果
|
||||
if valid_results:
|
||||
print("\n🏆 最高准确率的5个结果:")
|
||||
best_results = sorted(valid_results, key=lambda x: x["accuracy"], reverse=True)[
|
||||
:5
|
||||
]
|
||||
for i, r in enumerate(best_results, 1):
|
||||
print(
|
||||
f" {i}. 测试#{r['test_id']:3d}: 准确率{r['accuracy']:5.1f}% 长度{r['length']:2d} ({r['time_minutes']:4.1f}分钟)"
|
||||
)
|
||||
|
||||
print("\n⚠️ 最低准确率的5个结果:")
|
||||
worst_results = sorted(valid_results, key=lambda x: x["accuracy"])[:5]
|
||||
for i, r in enumerate(worst_results, 1):
|
||||
print(
|
||||
f" {i}. 测试#{r['test_id']:3d}: 准确率{r['accuracy']:5.1f}% 长度{r['length']:2d} ({r['time_minutes']:4.1f}分钟)"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
# 设置多进程启动方法（Linux上默认是fork，但spawn更安全）
|
||||
mp.set_start_method("spawn", force=True)
|
||||
main()
|
||||
except KeyboardInterrupt:
|
||||
print("\n\n⏹️ 测试被用户中断")
|
||||
# 清理可能的临时文件
|
||||
for f in Path(".").glob("test_fhe_output_*.jsonl"):
|
||||
f.unlink(missing_ok=True)
|
||||
except Exception as e:
|
||||
print(f"\n💥 程序异常: {e}")
|
||||
# 清理可能的临时文件
|
||||
for f in Path(".").glob("test_fhe_output_*.jsonl"):
|
||||
f.unlink(missing_ok=True)
|
||||
Reference in New Issue
Block a user