PyTorch 1.0尝鲜C++ API

PyTorch 1.0预览版发布了,其中一项重大更新就是支持C++调用了,对于工业产品发布来说是个极大利好。

怎么把用Python写的模型发布到C++,官方给出了答案:Loading a PyTorch Model in C++

这里把官方流程走了一遍,私心添加了OpenCV图像转Tensor的例子。

主要分2步:

  • 把PyTorch模型转成Torch Script,并保存
  • 在C++中调用Torch Script文件

先放代码: iamhankai/cpp-pytorch

1. PyTorch模型转Torch Script

官方提供了2种方式,Tracing和Annotation,一般来说,你的模型不依赖于输入的内容,用Tracing即可,因为它不需要改多少原代码,写几行代码就够了。

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

# save
traced_script_module.save("model.pt")

这样,你的PyTorch模型就序列化并保存在model.pt文件里了。

2. C++调用Torch Script

PyTorch官方发布的C++ API就是LibTorch,它提供了必要的库文件、头文件和CMake配置文件,现在已经可以在PyTorch首页下载了。

C++调用的一个例子如下

#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().   std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(argv[1]);

  assert(module != nullptr);
  std::cout << "ok\n";

  // Create a vector of inputs.   std::vector<torch::jit::IValue> inputs;
  inputs.push_back(torch::ones({1, 3, 224, 224}));

  // Execute the model and turn its output into a tensor.   auto output = module->forward(inputs).toTensor();

  std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
}

创建CMake工程,添加LibTorch和CMakeLists.txt,编译后就可以试用了。完整代码可参考 iamhankai/cpp-pytorch

完整的工程可以参考:

iamhankai/cpp-pytorchgithub.com《PyTorch 1.0尝鲜C++ API》

参考资料:

Torch Script – PyTorch master documentation

PyTorch C++ API – PyTorch master documentation

    原文作者:瑟木
    原文地址: https://zhuanlan.zhihu.com/p/45981883
    本文转自网络文章,转载此文章仅为分享知识,如有侵权,请联系博主进行删除。
点赞