TorchScript介绍及使用

 

TorchScript介绍

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

trace使用

Torchscript使用分为两个部分分别是script跟trace，其中trace是跟踪执行步骤，记录模型调用推理时执行的每个步骤，代码演示如下：

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


      def forward(self, x, h):
            new_h = torch.tanh(self.linear(x) + h)
            return new_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为例，代码演示如下：

import torch
import torchvision as tv

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()
with torch.no_grad():
    print(loaded_trace(list(im)))

 

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

#include  // One-stop header.
#include 
#include 
int main(int argc, const char* argv[]) {
    if (argc != 2) {
      std::cerr << "usage: example-app 
";
      return -1;
    }

    // Deserialize the ScriptModule from a file using torch::load().
    torch::Module module = torch::load(argv[1]);
    std::vector inputs;
    inputs.push_back(torch::randn({4, 8}));
    inputs.push_back(torch::randn({8, 5}));
    torch::Tensor output = module.forward(std::move(inputs)).toTensor();

    std::cout << output << std::endl;
}

 

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

 

审核编辑 ：李倩