198 lines
8.0 KiB
Python
198 lines
8.0 KiB
Python
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/embedding_20250930_102826/best_embedding_model.pth"
|
|
model_path = "../model/WholeIngredientRecognition/embedding_20251021_085915/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/WholeIngredientRecognition/embedding_20251021_085915/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() |