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)