TorchScript介紹及使用

TorchScript介紹

TorchScript是PyTorch模型推理部署的中間表示，可以在高性能環(huán)境libtorch（C ++）中直接加載，實(shí)現(xiàn)模型推理，而無需Pytorch訓(xùn)練框架依賴。torch.jit是torchscript Python語言包支持，支持pytorch模型快速，高效，無縫對接到libtorch運(yùn)行時(shí)，實(shí)現(xiàn)高效推理。它是Pytorch中除了訓(xùn)練部分之外，開發(fā)者最需要掌握的Pytorch框架開發(fā)技能之一。

trace使用

Torchscript使用分為兩個(gè)部分分別是script跟trace，其中trace是跟蹤執(zhí)行步驟，記錄模型調(diào)用推理時(shí)執(zhí)行的每個(gè)步驟，代碼演示如下：

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)

運(yùn)行結(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)))

最終得到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<std::endl;
}

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

審核編輯：李倩

聲明：本文內(nèi)容及配圖由入駐作者撰寫或者入駐合作網(wǎng)站授權(quán)轉(zhuǎn)載。文章觀點(diǎn)僅代表作者本人，不代表電子發(fā)燒友網(wǎng)立場。文章及其配圖僅供工程師學(xué)習(xí)之用，如有內(nèi)容侵權(quán)或者其他違規(guī)問題，請聯(lián)系本站處理。舉報(bào)投訴