轻松学Pytorch之torchscript使用！

点击上方蓝字关注我们

微信公众号：OpenCV学堂 关注获取更多计算机视觉与深度学习知识

TorchScript介绍

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

trace使用

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

代码语言：javascript复制

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)

运行结果如下：

代码语言：javascript复制

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

跟踪执行结果

代码语言：javascript复制

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)):
   : __torch__.torch.nn.modules.linear.Linear = prim::GetAttr[name="linear"](%self.1)
  ! : Tensor = prim::CallMethod[name="forward"](, %input)
   : int = prim::Constant[value=1]() # D:/python/pytorch_openvino_demo/ch5/faster_rcnn2onnx.py:112:0
   : Float(3:4, 4:1, requires_grad=1, device=cpu) = aten::add(!, %h, ) # D:/python/pytorch_openvino_demo/ch5/faster_rcnn2onnx.py:112:0
   : Float(3:4, 4:1, requires_grad=1, device=cpu) = aten::tanh() # D:/python/pytorch_openvino_demo/ch5/faster_rcnn2onnx.py:112:0
   : (Float(3:4, 4:1, requires_grad=1, device=cpu), Float(3:4, 4:1, requires_grad=1, device=cpu)) = prim::TupleConstruct(, )
  return ()

script部分使用

script是导出模型为中间IR格式文件，支持高性能libtorch C 部署，我们以torchvision中Mask-RCNN导出中间格式IR为例，代码演示如下：

代码语言：javascript复制

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 部分加载的代码如下：

代码语言：javascript复制

#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
int main(int argc, const char* argv[]) {
    if (argc != 2) {
      std::cerr << "usage: example-app <path-to-exported-script-module>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<torch::jit::IValue> 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系统化学习路线图