diff --git a/faiss_vector_db/build_faiss_index.py b/faiss_vector_db/build_faiss_index.py index 47f407d..6ee1059 100644 --- a/faiss_vector_db/build_faiss_index.py +++ b/faiss_vector_db/build_faiss_index.py @@ -267,8 +267,11 @@ class FAISSIndexBuilder: # 保存标签映射 labels_path = os.path.join(output_dir, 'labels.pkl') + labels_json_path = os.path.join(output_dir, 'labels.json') with open(labels_path, 'wb') as f: pickle.dump(self.labels, f) + with open(labels_json_path, "w", encoding="utf-8") as f: + json.dump(self.labels, f, ensure_ascii=False) print(f"标签映射已保存到: {labels_path}") # 保存类别信息 @@ -288,8 +291,16 @@ class FAISSIndexBuilder: # 保存特征向量（可选，用于调试） embeddings_path = os.path.join(output_dir, 'embeddings.npy') + embeddings_json_path = os.path.join(output_dir, 'embeddings.json') embeddings_array = np.array(self.embeddings) np.save(embeddings_path, embeddings_array) + + # 转成 Python 列表（JSON 不支持 numpy.array 直接序列化） + embeddings_list = embeddings_array.tolist() + # 保存成 JSON 文件 + with open(embeddings_json_path, "w", encoding="utf-8") as f: + json.dump(embeddings_list, f, ensure_ascii=False) + print(f"特征向量已保存到: {embeddings_path}") def build_complete_index(self, train_dir: str, output_dir: str, diff --git a/toAndroid/toAndroidEmbedding.py b/toAndroid/toAndroidEmbedding.py new file mode 100644 index 0000000..0a5ed58 --- /dev/null +++ b/toAndroid/toAndroidEmbedding.py @@ -0,0 +1,196 @@ +import torch +import torch.nn.functional as F +import os +import sys + +# 添加项目根目录到路径 +current_dir = os.path.dirname(os.path.abspath(__file__)) +parent_dir = os.path.dirname(current_dir) +sys.path.insert(0, parent_dir) + +from net.resnet_embedding import create_resnet50_embedding, create_mobile_resnet50_embedding + + +def main(): + print("开始转换ResNet50 Embedding模型为移动端格式...") + + # 1. 加载训练好的embedding模型权重 + base_model = create_resnet50_embedding(embedding_dim=512, pretrained=True) + model_path = "../model/embedding_20250917_145342/best_embedding_model.pth" + + if not os.path.exists(model_path): + print(f"错误：模型文件不存在 {model_path}") + return + + # 加载模型权重 + checkpoint = torch.load(model_path, map_location='cpu') + if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint: + # 如果保存的是完整的checkpoint + base_model.load_state_dict(checkpoint['model_state_dict']) + print("✓ 从checkpoint加载模型权重成功") + else: + # 如果保存的是纯模型权重 + base_model.load_state_dict(checkpoint) + print("✓ 基础模型权重加载成功") + + # 设置为评估模式 + base_model.eval() + + # 2. 创建移动端模型并复制权重 + # 直接使用基础模型，但在forward中包含预处理 + mobile_model = base_model # 使用同一个模型实例 + print("✓ 移动端模型创建成功") + + # 3. 测试模型 + # 创建示例输入 (模拟Android端输入：2x3x224x224，值范围0-1，使用batch_size=2避免BatchNorm问题) + example_input = torch.rand(2, 3, 224, 224, dtype=torch.float32) + + with torch.no_grad(): + # 测试基础模型 + base_output = base_model(example_input) + print(f"✓ 基础模型测试成功，输出形状: {base_output.shape}") + print(f" 基础模型输出值范围: [{base_output.min():.3f}, {base_output.max():.3f}]") + print(f" 基础模型输出L2范数: {torch.norm(base_output, p=2, dim=1)}") + + # 测试移动端模型（使用forward_mobile方法） + mobile_output = mobile_model.forward_mobile(example_input) + print(f"✓ 移动端模型测试成功，输出形状: {mobile_output.shape}") + print(f" 移动端模型输出值范围: [{mobile_output.min():.3f}, {mobile_output.max():.3f}]") + print(f" 移动端模型输出L2范数: {torch.norm(mobile_output, p=2, dim=1)}") + + # 检查输出是否一致（由于预处理不同，可能会有差异） + print("✓ 移动端模型测试完成（包含预处理）") + + # 4. 创建专门的移动端模型类用于TorchScript转换 + class MobileEmbeddingModel(torch.nn.Module): + def __init__(self, base_model): + super().__init__() + self.base_model = base_model + + def forward(self, x): + # 在forward中包含预处理 + return self.base_model.forward_mobile(x) + + # 创建移动端包装模型 + mobile_wrapper = MobileEmbeddingModel(base_model) + mobile_wrapper.eval() + + # 5. 转换为 TorchScript + try: + # 为TorchScript转换创建单个样本输入 (1x3x224x224) + single_input = torch.rand(1, 3, 224, 224, dtype=torch.float32) + + # 使用trace方法转换 + traced_model = torch.jit.trace(mobile_wrapper, single_input) + + # 保存模型 + output_path = "../model/embedding_20250917_145342/best_embedding_model_mobile.pt" + traced_model.save(output_path) + print(f"✓ TorchScript模型保存成功: {output_path}") + + # 验证保存的模型 + loaded_model = torch.jit.load(output_path) + with torch.no_grad(): + # 使用单个样本测试保存的模型 + loaded_output = loaded_model(single_input) + print(f"✓ 保存的模型验证成功，输出形状: {loaded_output.shape}") + print(f" 保存模型输出L2范数: {torch.norm(loaded_output, p=2, dim=1)}") + + # 用移动端模型处理单个样本进行比较 + mobile_single_output = mobile_wrapper(single_input) + + # 检查输出是否一致 + if torch.allclose(mobile_single_output, loaded_output, atol=1e-6): + print("✓ 模型输出一致性验证通过") + else: + print("⚠ 警告：模型输出存在差异") + print(f" 差异最大值: {torch.max(torch.abs(mobile_single_output - loaded_output))}") + + # 测试不同尺寸的输入 + print("\n测试不同输入尺寸...") + test_sizes = [(1, 3, 256, 256), (1, 3, 320, 320), (1, 3, 128, 128)] + + for size in test_sizes: + test_input = torch.rand(*size, dtype=torch.float32) + try: + with torch.no_grad(): + test_output = loaded_model(test_input) + print(f"✓ 输入尺寸 {size} 测试成功，输出形状: {test_output.shape}") + print(f" 输出L2范数: {torch.norm(test_output, p=2, dim=1)}") + except Exception as e: + print(f"✗ 输入尺寸 {size} 测试失败: {e}") + + except Exception as e: + print(f"✗ TorchScript转换失败: {e}") + import traceback + traceback.print_exc() + return + + # 5. 模型信息统计 + print("\n" + "="*60) + print("模型信息统计:") + print("="*60) + + # 计算模型大小 + model_size = os.path.getsize(output_path) / (1024 * 1024) # MB + print(f"模型文件大小: {model_size:.2f} MB") + + # 计算参数数量 + total_params = sum(p.numel() for p in base_model.parameters()) + trainable_params = sum(p.numel() for p in base_model.parameters() if p.requires_grad) + print(f"总参数数量: {total_params:,}") + print(f"可训练参数数量: {trainable_params:,}") + + print("\n" + "="*60) + print("移动端Embedding模型转换完成！") + print("="*60) + print("使用说明：") + print("1. Android端输入图片应为 [batch, 3, height, width] 格式") + print("2. 像素值范围：0-1 (float32)") + print("3. 模型会自动处理：") + print(" - 缩放到 224x224") + print(" - ImageNet标准化 (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])") + print("4. 输出：512维L2归一化的embedding向量") + print("5. 可用于：") + print(" - 图片相似度计算 (余弦相似度)") + print(" - 图片检索和匹配") + print(" - 特征向量数据库构建") + print("="*60) + + # 6. 生成示例代码 + print("\nAndroid端使用示例 (Java/Kotlin):") + print("```java") + print("// 加载模型") + print("Module module = LiteModuleLoader.load(assetFilePath(\"best_embedding_model_mobile.pt\"));") + print("") + print("// 准备输入 (假设bitmap已转换为float数组，值范围0-1)") + print("float[] inputArray = ...; // [1, 3, height, width]") + print("Tensor inputTensor = Tensor.fromBlob(inputArray, new long[]{1, 3, height, width});") + print("") + print("// 推理") + print("IValue[] outputs = module.forward(IValue.from(inputTensor)).toTuple();") + print("Tensor outputTensor = outputs[0].toTensor();") + print("float[] embedding = outputTensor.getDataAsFloatArray(); // 512维向量") + print("") + print("// 计算相似度") + print("float similarity = cosineSimilarity(embedding1, embedding2);") + print("```") + + +def test_model_conversion(): + """ + 测试模型转换的完整流程 + """ + print("开始模型转换测试...") + + try: + main() + print("✓ 模型转换测试完成") + except Exception as e: + print(f"✗ 模型转换测试失败: {e}") + import traceback + traceback.print_exc() + + +if __name__ == "__main__": + main() \ No newline at end of file