import torch
import torch.nn as nn
import torch.nn.functional as F
from net import create_food_cnn


# 1. 初始化模型
model = create_food_cnn()
# 2. 加载训练好的权重
model.load_state_dict(torch.load("../model/02/best_food_model.pth", map_location='cpu'))
model.eval()  # 设置为推理模式

# 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/02/best_food_model_mobile.pt")