轻松学Pytorch之torchscript使用！

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.py:112:0
  %13 : Float(3:4, 4:1, requires_grad=1, device=cpu) = aten::add(%21, %h, %12) # D:/python/pytorch_openvino_demo/ch5/faster_rcnn2onnx.py:112:0
  %14 : Float(3:4, 4:1, requires_grad=1, device=cpu) = aten::tanh(%13) # D:/python/pytorch_openvino_demo/ch5/faster_rcnn2onnx.py:112:0
  %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是导出模型为中间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)))

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

    // Deserialize the ScriptModule from a file using torch::jit::load().
    torch::jit::script::Module module = torch::jit::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;
}

扫码查看OpenCV+OpenVIO+Pytorch系统化学习路线图