美文网首页
Pytorch_1: libtorch的C++调用(ubuntu

Pytorch_1: libtorch的C++调用(ubuntu

作者: 闪电侠悟空 | 来源:发表于2019-01-08 14:50 被阅读0次

    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.经验

    相关文章

      网友评论

          本文标题:Pytorch_1: libtorch的C++调用(ubuntu

          本文链接:https://www.haomeiwen.com/subject/fkivrqtx.html