菜品识别程序,增加向量输出到控制台的功能!
This commit is contained in:
@@ -7,7 +7,7 @@ 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.load_state_dict(torch.load("../model/06/best_food_model.pth", map_location='cpu'))
|
||||
model.eval() # 设置为推理模式
|
||||
|
||||
# 3. 创建示例输入 (假设输入是 3x224x224 的图片)
|
||||
@@ -15,4 +15,4 @@ 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")
|
||||
traced_script_module.save("../model/06/best_food_model_mobile.pt")
|
||||
|
||||
Reference in New Issue
Block a user