yixiong/FoodClassifier
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修改转安卓模型的代码，增加加载cosFace模型的逻辑。

Browse Source
This commit is contained in:
zhangpu
2025-11-06 11:01:06 +08:00
parent 2c74c81918
commit 3ca5d6e7d0
1 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
toAndroid/toAndroidEmbedding.py
+10 -4
View File
@@ -17,8 +17,8 @@ def main():
# 1. 加载训练好的embedding模型权重 # 1. 加载训练好的embedding模型权重
# base_model = create_resnet50_embedding(embedding_dim=512, pretrained=True) # base_model = create_resnet50_embedding(embedding_dim=512, pretrained=True)
base_model = create_mobile_resnet50_embedding(embedding_dim=512, pretrained=True) base_model = create_mobile_resnet50_embedding(embedding_dim=512, pretrained=True)
# model_path = "../model/embedding_20250930_102826/best_embedding_model.pth" model_path = "../model/DishClassification/cosface_20251105_200551/best_cosface_model.pth"
model_path = "../model/WholeIngredientRecognition/embedding_20251030_135024/best_embedding_model.pth" # model_path = "../model/WholeIngredientRecognition/embedding_20251030_135024/best_embedding_model.pth"
# model_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model.pth" # model_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model.pth"
if not os.path.exists(model_path): if not os.path.exists(model_path):
@@ -30,7 +30,12 @@ def main():
if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint: if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
# 如果保存的是完整的checkpoint # 如果保存的是完整的checkpoint
base_model.load_state_dict(checkpoint['model_state_dict']) base_model.load_state_dict(checkpoint['model_state_dict'])
print("✓ 从checkpoint加载模型权重成功") 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: else:
# 如果保存的是纯模型权重 # 如果保存的是纯模型权重
base_model.load_state_dict(checkpoint) base_model.load_state_dict(checkpoint)
@@ -87,7 +92,8 @@ def main():
traced_model = torch.jit.trace(mobile_wrapper, single_input) traced_model = torch.jit.trace(mobile_wrapper, single_input)
# 保存模型 # 保存模型
output_path = "../model/WholeIngredientRecognition/embedding_20251030_135024/best_embedding_model_mobile.pt" output_path = "../model/DishClassification/cosface_20251105_200551/best_embedding_model_mobile.pt"
# output_path = "../model/WholeIngredientRecognition/embedding_20251030_135024/best_embedding_model_mobile.pt"
# output_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model_mobile.pt" # output_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model_mobile.pt"
traced_model.save(output_path) traced_model.save(output_path)
print(f"✓ TorchScript模型保存成功: {output_path}") print(f"✓ TorchScript模型保存成功: {output_path}")