pytorch模型转成pt(c++可调用)

参考：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')