223 lines
8.3 KiB
Python
223 lines
8.3 KiB
Python
import os
|
||
import sys
|
||
import json
|
||
import math
|
||
from typing import List, Tuple, Union
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
# 添加项目根目录到路径（与 toAndroidEmbedding.py 保持一致）
|
||
# current_dir = os.path.dirname(os.path.abspath(__file__))
|
||
# parent_dir = os.path.dirname(current_dir)
|
||
# sys.path.insert(0, parent_dir)
|
||
print(sys.path)
|
||
|
||
|
||
def _load_embeddings_json(json_path: str) -> torch.Tensor:
|
||
"""
|
||
读取 ../faiss_vector_db/embeddings.json，尽量兼容常见格式，并返回 [N, 512] 的 float32 Tensor。
|
||
兼容格式示例：
|
||
1) [ [float]*512, [float]*512, ... ]
|
||
2) [ {"embedding": [float]*512, ...}, {"embedding": [...]}, ... ]
|
||
3) {"embeddings": [ [float]*512, ... ] }
|
||
4) {"items": [ {"embedding": [...]}, ... ] }
|
||
|
||
如维度不为512或为空，抛出异常。
|
||
"""
|
||
if not os.path.exists(json_path):
|
||
raise FileNotFoundError(f"未找到 embeddings.json: {json_path}")
|
||
|
||
with open(json_path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
def extract_vectors(obj: Union[list, dict]) -> List[List[float]]:
|
||
if isinstance(obj, list):
|
||
# list of vectors or list of dicts
|
||
if len(obj) == 0:
|
||
return []
|
||
if isinstance(obj[0], list):
|
||
return obj
|
||
if isinstance(obj[0], dict):
|
||
vecs = []
|
||
for item in obj:
|
||
if "embedding" in item and isinstance(item["embedding"], list):
|
||
vecs.append(item["embedding"])
|
||
return vecs
|
||
return []
|
||
elif isinstance(obj, dict):
|
||
# possible keys: "embeddings", "items", "data"
|
||
for key in ("embeddings", "items", "data", "vectors"):
|
||
if key in obj:
|
||
inner = obj[key]
|
||
return extract_vectors(inner)
|
||
# maybe dict of id->vector
|
||
# ensure values are lists
|
||
values = list(obj.values())
|
||
if values and isinstance(values[0], list):
|
||
return values
|
||
if values and isinstance(values[0], dict) and "embedding" in values[0]:
|
||
return [v["embedding"] for v in values if "embedding" in v]
|
||
return []
|
||
else:
|
||
return []
|
||
|
||
vectors = extract_vectors(data)
|
||
if not vectors:
|
||
raise ValueError("embeddings.json 未解析到任何向量，请确认文件结构。")
|
||
|
||
# 验证维度
|
||
dim = len(vectors[0])
|
||
if dim != 512:
|
||
raise ValueError(f"向量维度为 {dim}，与期望的512不符。请检查 embeddings.json。")
|
||
|
||
# 转为 Tensor
|
||
emb_tensor = torch.tensor(vectors, dtype=torch.float32) # [N, 512]
|
||
if emb_tensor.ndim != 2 or emb_tensor.shape[1] != 512:
|
||
raise ValueError("向量张量形状异常，期望 [N, 512]。")
|
||
|
||
return emb_tensor
|
||
|
||
|
||
class TwoLayerSimilarityNet(nn.Module):
|
||
"""
|
||
两层网络：
|
||
- layer1: Linear(512->512) 无偏置，初始化为恒等（可视作占位层，满足“两层”要求）
|
||
- layer2: Linear(512->N) 无偏置，权重为数据库向量（每行一个向量），用于点积相似度
|
||
|
||
为近似余弦相似度：
|
||
- 构造时将数据库向量 L2 归一化后作为 layer2.weight
|
||
- 前向时对输入向量做 L2 归一化，再依次通过 layer1、layer2
|
||
"""
|
||
|
||
def __init__(self, db_vectors: torch.Tensor):
|
||
"""
|
||
db_vectors: [N, 512] float32
|
||
"""
|
||
super().__init__()
|
||
|
||
if db_vectors.ndim != 2 or db_vectors.shape[1] != 512:
|
||
raise ValueError("db_vectors 形状必须为 [N, 512]")
|
||
|
||
N = db_vectors.shape[0]
|
||
|
||
# （已移除恒等层，直接使用单层相似度投影）
|
||
|
||
# 单层：权重为数据库向量（行表示样本），无偏置
|
||
self.layer2 = nn.Linear(512, N, bias=False)
|
||
|
||
# 归一化数据库向量 -> 近似余弦相似度
|
||
# db_norm = F.normalize(db_vectors, p=2, dim=1) # [N, 512], 每行单位范数
|
||
with torch.no_grad():
|
||
# nn.Linear 的 weight 形状为 [out_features, in_features] = [N, 512]
|
||
# self.layer2.weight.copy_(db_norm)
|
||
self.layer2.weight.copy_(db_vectors)
|
||
|
||
# 冻结所有参数（移动端推理仅用作检索头）
|
||
for p in self.parameters():
|
||
p.requires_grad = False
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
x: [B, 512] float32
|
||
返回: [B, N] 相似度（点积，因归一化近似余弦）
|
||
"""
|
||
if x.ndim != 2 or x.shape[1] != 512:
|
||
raise ValueError("输入张量形状必须为 [batch, 512]")
|
||
|
||
# 输入归一化
|
||
x_norm = F.normalize(x, p=2, dim=1) # [B, 512]
|
||
|
||
# 单层线性变换
|
||
sims = self.layer2(x_norm) # [B, N]
|
||
# sims = self.layer2(x) # [B, N]
|
||
|
||
return sims
|
||
|
||
|
||
def build_and_export_similarity_head(
|
||
embeddings_json_path: str = "faiss_vector_db/embeddings.json",
|
||
output_path: str = "faiss_vector_db/similarity_head.pt",
|
||
) -> Tuple[str, int]:
|
||
"""
|
||
构建两层相似度网络并导出 TorchScript。
|
||
返回 (输出路径, 样本数)
|
||
"""
|
||
print("开始构建相似度检索头（TwoLayerSimilarityNet）...")
|
||
|
||
# 1) 加载原始向量
|
||
db_vectors = _load_embeddings_json(embeddings_json_path) # [N, 512]
|
||
N = db_vectors.shape[0]
|
||
print(f"✓ 已加载数据库向量，样本数: {N}，维度: 512")
|
||
|
||
# 2) 构建网络
|
||
model = TwoLayerSimilarityNet(db_vectors)
|
||
model.eval()
|
||
print("✓ 单层检索头构建完成（layer2=向量库权重，无偏置）")
|
||
|
||
# 3) 进行一次简单测试
|
||
with torch.no_grad():
|
||
test_input = torch.rand(2, 512, dtype=torch.float32)
|
||
test_output = model(test_input)
|
||
print(f"✓ 前向测试成功，输入形状: {tuple(test_input.shape)}，输出形状: {tuple(test_output.shape)}")
|
||
print(f" 输出值范围: [{test_output.min().item():.4f}, {test_output.max().item():.4f}]")
|
||
|
||
# 4) 导出为 TorchScript（使用 trace，输入为 [1, 512]）
|
||
try:
|
||
example = torch.rand(1, 512, dtype=torch.float32)
|
||
traced = torch.jit.trace(model, example)
|
||
traced.save(output_path)
|
||
print(f"✓ TorchScript 模型已保存到: {output_path}")
|
||
|
||
# 验证加载
|
||
loaded = torch.jit.load(output_path)
|
||
with torch.no_grad():
|
||
out = loaded(example)
|
||
if out.shape != (1, N):
|
||
raise RuntimeError("加载后的 TorchScript 输出形状异常")
|
||
print("✓ 导出模型验证通过")
|
||
except Exception as e:
|
||
print(f"✗ TorchScript 导出失败: {e}")
|
||
raise
|
||
|
||
# 5) 信息统计
|
||
model_size_mb = os.path.getsize(output_path) / (1024 * 1024.0)
|
||
total_params = sum(p.numel() for p in model.parameters())
|
||
print("\n" + "=" * 60)
|
||
print("检索头模型信息:")
|
||
print("=" * 60)
|
||
print(f"样本数 (输出维度): {N}")
|
||
print(f"模型文件大小: {model_size_mb:.2f} MB")
|
||
print(f"参数总数: {total_params:,} (全部冻结)")
|
||
print("=" * 60)
|
||
|
||
return output_path, N
|
||
|
||
|
||
def main():
|
||
print("开始转换检索头为移动端格式（TorchScript）...")
|
||
embeddings_path = "../faiss_vector_db/DishClassification/faiss_index/embeddings.json"
|
||
output_path = "../faiss_vector_db/DishClassification/faiss_index/similarity_head.pt"
|
||
|
||
try:
|
||
path, n_items = build_and_export_similarity_head(embeddings_path, output_path)
|
||
print("\nAndroid端使用说明：")
|
||
print("1) 输入为 [batch, 512] 的 float32 张量（由移动端Embedding模型输出）")
|
||
print("2) 输出为 [batch, N] 的相似度，其中 N 为向量库样本数")
|
||
print("3) 打包到APK后，使用 PyTorch Mobile 加载 similarity_head.pt 进行推理")
|
||
print("4) 如需做 Top-K 检索，取每行的最大/Top-K即可")
|
||
print("\n示例（Kotlin/Java 伪代码）：")
|
||
print(" val module = LiteModuleLoader.load(assetFilePath(\"similarity_head.pt\"))")
|
||
print(" val inputTensor = Tensor.fromBlob(inputArray, longArrayOf(batch.toLong(), 512L))")
|
||
print(" val outputTensor = module.forward(IValue.from(inputTensor)).toTensor()")
|
||
print(" val sims = outputTensor.dataAsFloatArray // 长度 = batch * N")
|
||
print("\n完成。")
|
||
except Exception as e:
|
||
print(f"✗ 转换失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
|
||
|
||
main() |