Pytorch_1: libtorch的C++调用(ubuntu

1.获取与安装

• 在主页面选择libtorch 下载
• 解压得到libtorch包(包括bin,include,lib,share,build-version)
• 安装就是把libtorch包放置到自己想要的地方(就放置在下载处也ok)

2.使用

• CMakelists.txt

project(test_pytorch)
set(CMAKE_CXX_STANDARD 11)
set(Torch_DIR /home/xxx/下载/libtorch/share/cmake/Torch) #指定libTorch位置(应该是有更好的办法安装)

#include_directories(${OpenCV_INCLUDE_DIRS} /home/xxx/下载/libtorch/include /home/xxx/下载/libtorch/include/torch/csrc/api/include)
find_package(OpenCV REQUIRED)
find_package(Torch REQUIRED)    # 自动查找libTorch包

add_executable(test_pytorch test_pytorch.cpp)
target_link_libraries(test_pytorch ${OpenCV_LIBS}  ${TORCH_LIBRARIES}) # 加入libTorch的库文件路径

• C++文件的编写

#include <iostream>
#include "torch/script.h"
#include "torch/torch.h"
using namespace std;

int main() {
    //加载pytorch模型
    std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(
            "/media/xxx/model.pt");
    assert(module != nullptr);

    //加载输入
    // 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.
    at::Tensor output = module->forward(inputs).toTensor();
    std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
    cout << "Hello, world!" << endl;
}

3.选择计算Device(CPU或GPU)
最小代码

//设置Device
    torch::DeviceType device_type; //设置Device类型
    device_type = torch::kCUDA;  //torch::kCUDA  and torch::kCPU
    torch::Device device(device_type, 0);
//把模型和数据都送到Device中去(数据和模型必须都在同一个device,结果也是)
    module->to(device);
    inputs.push_back(torch::ones({1, 3, 224, 224}).to(device));

完整代码:

#include <iostream>
#include "torch/script.h"
#include "torch/torch.h"
using namespace std;

int main() {
    //加载pytorch模型
    std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(
            "/media/xxxx/model.pt");//注意检查路径
    assert(module != nullptr);

    torch::DeviceType device_type; //设置Device类型
    device_type = torch::kCUDA;  //torch::kCUDA  and torch::kCPU
    torch::Device device(device_type, 0);

    //模型转到GPU中去
    module->to(device);
    torch::cuda::is_available(); //判断是否支持GPU加速
    // Create a vector of inputs.
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::ones({1, 3, 224, 224}).to(device));

    //指定执行位置
    // Execute the model and turn its output into a tensor.
    for (int i = 0; i < 100; i++) {
        at::Tensor output = module->forward(inputs).toTensor();
        std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
    }
}

5.经验

• 学习新框架的最好方式就是查找源代码,最后一个Github的源代码帮助很大.
  4.参考资料
• 官网说明 LOADING A PYTORCH MODEL IN C++
• 利用Pytorch的C++前端(libtorch)读取预训练权重并进行预测
• a github example