19 lines
568 B
Python
19 lines
568 B
Python
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/06/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/06/best_food_model_mobile.pt")
|