TorchScript介绍及使用

TorchScript介绍

TorchScript是PyTorch模型推理部署的中间表示，可以在高性能环境libtorch（C ++）中直接加载，实现模型推理，而无需Pytorch训练框架依赖。torch.jit是torchscript Python语言包支持，支持pytorch模型快速，高效，无缝对接到libtorch运行时，实现高效推理。它是Pytorch中除了训练部分之外，开发者最需要掌握的Pytorch框架开发技能之一。

trace使用

Torchscript使用分为两个部分分别是script跟trace，其中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)

运行结果如下：

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

跟踪执行结果

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是导出模型为中间IR格式文件，支持高性能libtorch C++部署，我们以torchvision中Mask-RCNN导出中间格式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)))

最终得到torchscript文件，支持直接通过libtorch部署，其中通过torchscript C++部分加载的代码如下：

#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<< output << std::endl;
}

上面代码来自官方测试程序，特别说明一下！

审核编辑 ：李倩