无法在PyTorch Mobile中加载PyTorch模型的问题排查

问题根源 #

你直接用torch.save(model, path)保存的是Python环境下的模型序列化对象，包含Python特有的依赖结构，而PyTorch Mobile的C++ runtime只能解析TorchScript格式的模型，所以会出现找不到archive/constants.pkl的错误。

解决步骤 #

1. 将模型转换为TorchScript格式

PyTorch提供两种转换方式，根据模型结构选择：

• 追踪（Trace）方式：适合没有动态控制流（如if/else、循环）的模型

  import torch
  from your_model_module import YourModel  # 替换为你的模型类导入
  
  # 从state dict恢复模型
  model = YourModel()
  model.load_state_dict(torch.load("your_state_dict.pth"))
  model.eval()  # 必须切换到评估模式
  
  # 创建与实际输入匹配的示例张量（shape、dtype要一致）
  example_input = torch.randn(1, 3, 224, 224)  # 示例：图像分类模型输入
  # 转换并保存TorchScript模型
  traced_model = torch.jit.trace(model, example_input)
  traced_model.save("fullap.pt")
  

• 脚本（Script）方式：适合包含动态控制流的模型

  import torch
  from your_model_module import YourModel
  
  model = YourModel()
  model.load_state_dict(torch.load("your_state_dict.pth"))
  model.eval()
  
  scripted_model = torch.jit.script(model)
  scripted_model.save("fullap.pt")
  

2. 验证转换后的模型

在PC端先测试加载，确保模型可用：

loaded_model = torch.jit.load("fullap.pt")
output = loaded_model(example_input)  # 用示例输入测试推理

3. 重新部署到Flutter项目

• 替换原assets/models/fullap.pt为转换后的模型文件
• 确认pubspec.yaml中assets配置正确：

  flutter:
    assets:
      - assets/models/fullap.pt
  

关键注意事项 #

• 转换前必须调用model.eval()，避免BatchNorm、Dropout等层的训练行为干扰
• 示例输入的维度、数据类型必须和实际部署时的输入完全匹配
• 永远不要用torch.save直接保存模型给PyTorch Mobile使用，必须通过TorchScript转换

内容的提问来源于stack exchange，提问作者Manav Shah