TorchScript介紹及使用

TorchScript介紹

trace使用

classMyCell(torch.nn.Module):
def__init__(self):
super(MyCell,self).__init__()
self.linear=torch.nn.Linear(4,4)


defforward(self,x,h):
new_h=torch.tanh(self.linear(x)+h)
returnnew_h,new_h


my_cell=MyCell()
x,h=torch.rand(3,4),torch.rand(3,4)
traced_cell=torch.jit.trace(my_cell,(x,h))
print(traced_cell)
traced_cell(x,h)
print(traced_cell.graph)

運行結(jié)果如下：

MyCell(
original_name=MyCell
(linear):Linear(original_name=Linear)
)

跟蹤執(zhí)行結(jié)果

graph(%self.1:__torch__.MyCell,
%input:Float(3:4,4:1,requires_grad=0,device=cpu),
%h:Float(3:4,4:1,requires_grad=0,device=cpu)):
%19:__torch__.torch.nn.modules.linear.Linear=prim::GetAttr[name="linear"](%self.1)
%21:Tensor=prim::CallMethod[name="forward"](%19,%input)
%12:int=prim::Constant[value=1]()#D:/python/pytorch_openvino_demo/ch5/faster_rcnn2onnx.py0
%13:Float(3:4,4:1,requires_grad=1,device=cpu)=aten::add(%21,%h,%12)#D:/python/pytorch_openvino_demo/ch5/faster_rcnn2onnx.py0
%14:Float(3:4,4:1,requires_grad=1,device=cpu)=aten::tanh(%13)#D:/python/pytorch_openvino_demo/ch5/faster_rcnn2onnx.py0
%15:(Float(3:4,4:1,requires_grad=1,device=cpu),Float(3:4,4:1,requires_grad=1,device=cpu))=prim::TupleConstruct(%14,%14)
return(%15)

script部分使用

script是導(dǎo)出模型為中間IR格式文件，支持高性能libtorch C++部署，我們以torchvision中Mask-RCNN導(dǎo)出中間格式IR為例，代碼演示如下：

importtorch
importtorchvisionastv

num_classes=3
model=tv.models.detection.maskrcnn_resnet50_fpn(
pretrained=False,progress=True,
num_classes=num_classes,
pretrained_backbone=True)
im=torch.zeros(1,3,*(1333,800)).to("cpu")
model.load_state_dict(torch.load("D:/gaobao_model.pth"))
model=model.to("cpu")
model.eval()
ts=torch.jit.script(model)
ts.save("gaobao.ts")

loaded_trace=torch.jit.load("gaobao.ts")
loaded_trace.eval()
withtorch.no_grad():
print(loaded_trace(list(im)))

#include//One-stopheader.
#include
#include
intmain(intargc,constchar*argv[]){
if(argc!=2){
std::cerr<"usage:example-app
";
return-1;
}

//DeserializetheScriptModulefromafileusingtorch::load().
torch::Modulemodule=torch::load(argv[1]);
std::vectorinputs;
inputs.push_back(torch::randn({4,8}));
inputs.push_back(torch::randn({8,5}));
torch::Tensoroutput=module.forward(std::move(inputs)).toTensor();

std::cout<std::endl;
}

上面代碼來自官方測試程序，特別說明一下！

審核編輯 ：李倩