极智开发 | ubuntu 安装 libtorch

Posted 2022-10-31 极智视界

  欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享

  大家好，我是极智视界，本文介绍一下 ubuntu 安装 libtorch。

  libtorch 是 pytorch 的 C++ 版本，支持 CPU 端和 GPU 端的模型部署。相较于用 tensorrt 部署 pytorch 模型，用 libtorch的优势在于：pytorch 和 libtorch 同属一个生态，API 语句比较接近，并且不会出现某网络层不支持的问题。这里咱们来看一下 libtorch 在 ubuntu 上是怎么安装的。

  下面开始。

  首先需要到 pytorch 的官网去下载 libtorch 包，附官网地址：https://pytorch.org/

  这里咱们选择 GPU 版本的 libtorch，获取到 url 后，回到 ubuntu 中使用如下命令进行下载：

# 下载包
wget https://download.pytorch.org/libtorch/lts/1.8/cu111/libtorch-cxx11-abi-shared-with-deps-1.8.2%2Bcu111.zip

# 解压包
unzip libtorch-cxx11-abi-shared-with-deps-1.8.2%2Bcu111.zip

# 会得到一个 libtorch 包文件夹
|- libtorch
|  - include
|    - ...
|  - lib
|    - ...
|  - bin
|    - ...
|  - share
|    - ...
|  - build-hash
|  - build_version

  接着来编写一个简单工程进行调用测试。

mkdir demo

# 编写 CMakeList.txt
vim CMakeList.txt

  在 CMakeList.txt 内键入以下内容：

cmake_minimum_required(VERSION 2.8 FATAL_ERROR)
find_package(PythonInterp REQUIRED)

project(demo)

set(Torch_DIR ~/libtorch/share/cmake/Torch)  #你解压的libtorch的绝对路径
find_package(Torch REQUIRED)

set(CMAKE_CXX_FLAGS "$CAMKE_CXX_FLAGS $TORCH_CXX_FLAGS")
set(CUDA_INCLUDE_DIRS "/usr/local/cuda/include")

#main.cpp exe
add_executable(demo demo.cpp)

#link libtorch .a .so
target_link_libraries(demo "$TORCH_LIBRARIES")
target_include_directories(demo PRIVATE CUDA_INCLUDE_DIRS)

set_property(TARGET demo PROPERTY CXX_STANDARD 14)

  接着编写 demo.cpp，就简单打印一个 torch:ones(3) 3维都是1的张量:

// demo.cpp
#include<torch/torch.h>
#include<iostream>
//using namespace std;

int main()
    torch::Tensor tensor = torch::ones(3);
    std::cout << tensor << std::endl;

  开始编译 & 执行：

mkdir build
cd build
cmake ..
make -j8

# 生成 demo 可执行程序
# 执行 demo
./demo

  正常执行输出：

  好了，以上分享了 ubuntu 安装 libtorch 的方法。希望我的分享能对你的学习有一点帮助。