参考:pytorch官网
pytorch模型转化成C++可调用的序列化文件,其步骤如下(用trace的方法):
1. 加载模型参数,并设置成eval模式。
2. 定义输入,与训练模型时的输入大小相同。
3. 进行一次前向传播,记录下trace, 保存pt文件。
import torch
from MobileNetV2 import mobilenet_v2
model = mobilenet_v2()
state_dict = torch.load("checkpoint/mobilenet-v2_0.pth")
model.load_state_dict(state_dict)
model.eval()
x = torch.rand(1,3,128,128)
ts = torch.jit.trace(model, x)
ts.save('mobilenet_v2.pt')