pytorch 离线保存加载模型 torch.jit

1.保存模型

注意：torch.jit 保存模型，需要注意的是模型的forward方法不应该有判断语句，因为torch.jit不能识别判断语句，如果有请将其以其他形式进行替换；

1）. x = torch.rand(1,3,608,608).cuda()   #占位符

2）.trace_model = torch.jit.trace(model,x) 

3)  .torch.jit.save(trace_model,save_path)

通过字典序列化tensor，无需模型结构，不依赖于Model class，导入模型后即可使用；

2.加载模型

import torch
import io

path = 'path/your/best.pth'
model = torch.jit.load(path, map_location=torch.device('cuda'))

x = torch.rand(10,3,608,608).cuda() # this is the batch tensors of detect work
pre = model(x)
print(pre.shape)