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) base_model = create_mobile_resnet50_embedding(embedding_dim=512, pretrained=True) # model_path = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth" # model_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_cosface_model.pth" model_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth" # model_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/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("✓ 检测到Triplet模型格式，使用'model_state_dict'加载") elif 'backbone_state_dict' in checkpoint: # CosFace格式：使用backbone_state_dict（只加载backbone部分） base_model.load_state_dict(checkpoint['backbone_state_dict']) print("✓ 检测到CosFace模型格式，使用'backbone_state_dict'加载") 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/DishClassification/grid_search_20251121_102723/best_embedding_model_mobile.pt" output_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/best_embedding_model_mobile.pt" # output_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_embedding_model_mobile.pt" # output_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/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()