在模型定义中,增加了移动端的前向传播过程。修改了toAndroid的代码。

This commit is contained in:
zhanghuan
2025-09-11 14:35:46 +08:00
parent 4cfbfecf9d
commit 2f6ce5ffc3
3 changed files with 142 additions and 15 deletions
+70 -13
View File
@@ -1,18 +1,75 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from net import create_food_cnn
from net import create_food_cnn, create_mobile_food_cnn
import os
# 1. 初始化模型
model = create_food_cnn()
# 2. 加载训练好的权重
model.load_state_dict(torch.load("../model/06/best_food_model.pth", map_location='cpu'))
model.eval() # 设置为推理模式
def main():
print("开始转换模型为移动端格式...")
# 1. 加载训练好的基础模型权重
base_model = create_food_cnn()
model_path = "../model/07/best_food_model.pth"
if not os.path.exists(model_path):
print(f"错误:模型文件不存在 {model_path}")
return
base_model.load_state_dict(torch.load(model_path, map_location='cpu'))
print("✓ 基础模型权重加载成功")
# 2. 创建移动端模型并复制权重
mobile_model = create_mobile_food_cnn()
mobile_model.load_state_dict(base_model.state_dict())
mobile_model.eval()
print("✓ 移动端模型创建成功")
# 3. 测试模型
# 创建示例输入 (模拟Android端输入:1x3x224x224,值范围0-255)
example_input = torch.randint(0, 256, (1, 3, 224, 224), dtype=torch.float32)
with torch.no_grad():
output = mobile_model(example_input)
print(f"✓ 模型测试成功,输出形状: {output.shape}")
print(f" 输出值范围: [{output.min():.3f}, {output.max():.3f}]")
# 4. 转换为 TorchScript
try:
traced_model = torch.jit.trace(mobile_model, example_input)
# 保存模型
output_path = "../model/07/best_food_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(example_input)
print(f"✓ 保存的模型验证成功,输出形状: {loaded_output.shape}")
# 检查输出是否一致
if torch.allclose(output, loaded_output, atol=1e-6):
print("✓ 模型输出一致性验证通过")
else:
print("⚠ 警告:模型输出存在差异")
except Exception as e:
print(f"✗ TorchScript转换失败: {e}")
return
print("\n" + "="*50)
print("移动端模型转换完成!")
print("="*50)
print("使用说明:")
print("1. Android端输入图片应为 [batch, 3, height, width] 格式")
print("2. 像素值范围:0-255 (uint8 或 float32)")
print("3. 模型会自动处理:")
print(" - 归一化到 [0,1]")
print(" - 缩放到 32x32")
print(" - ImageNet标准化")
print("4. 输出:3个类别的logits")
print("="*50)
# 3. 创建示例输入 (假设输入是 3x224x224 的图片)
example_input = torch.randn(1, 3, 224, 224)
# 4. 转换为 TorchScript
traced_script_module = torch.jit.trace(model, example_input)
traced_script_module.save("../model/06/best_food_model_mobile.pt")
if __name__ == "__main__":
main()