当前位置：首页 > article >正文

ubuntu配置libtorch CPU版本

article 2024/9/24 15:25:01

配置环境：Ubuntu 20.04
Date：2024 / 08

1、下载最新版本的libtorch

wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip
unzip libtorch-shared-with-deps-latest.zip

2、创建一个C++工程文件夹，目录结构如下：

example-app/
  CMakeLists.txt
  example-app.cpp

其中，CMakeLists.txt文件如下：

cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
project(example-app)

# set(CMAKE_PREFIX_PATH "/path/to/libtorch;/another/path") #设置多个路径
set(CMAKE_PREFIX_PATH "/path/to/libtorch")

find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 17)

# The following code block is suggested to be used on Windows.
# According to https://github.com/pytorch/pytorch/issues/25457,
# the DLLs need to be copied to avoid memory errors.
if (MSVC)
  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
  add_custom_command(TARGET example-app
                     POST_BUILD
                     COMMAND ${CMAKE_COMMAND} -E copy_if_different
                     ${TORCH_DLLS}
                     $<TARGET_FILE_DIR:example-app>)
endif (MSVC)

example-app.cpp文件如下：

#include <torch/torch.h>
#include <iostream>

int main() {
  torch::Tensor tensor = torch::rand({2, 3});
  std::cout << tensor << std::endl;
}

3、编译
在example-app文件夹下打开终端

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/absolute/path/to/libtorch .. #这个路径可以直接在cmakelist中定义
cmake --build . --config Release

4、运行测试

./example-app

可以打印出数组证明安装成功

5、加载.pt文件cpp使用示例：

#include <iostream>
#include <cmath>

#include <torch/torch.h>      // Libtorch头文件
#include <torch/script.h>     // 需要加载PyTorch脚本模型
	

// libtorch模型加载和推理函数
torch::Tensor load_model_and_predict(torch::Tensor input_tensor) {
    // 加载已保存的PyTorch脚本模型（torch::jit::script::Module）
    torch::jit::script::Module module;
    try {
        module = torch::jit::load("/path/to/policy.pt");  // 模型路径
    }
    catch (const c10::Error& e) {
        std::cerr << "Error loading the model\n";
        return torch::Tensor();  // 返回一个空的张量，表示加载失败
    }

    // 模型推理
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(input_tensor);
    at::Tensor output = module.forward(inputs).toTensor();

    return output;
}

int main(int argc, char** argv) {

    // 构造输入张量进行模型推理（假设输入为1维张量，大小为48*15）
    torch::Tensor input_tensor = torch::rand({1, 48*15});//根据自己打模型输入设置
    torch::Tensor prediction = load_model_and_predict(input_tensor);

    std::cout << "Model prediction: " << prediction << std::endl;
    //std::cout << "input_tensor: " << input_tensor << std::endl;

    return 0;
}