Linux环境下在CPP项目内引用torch相关库完成编译

前言与需求

Pytorch应该是现在最流行的ai框架之一了，之前在学校日常使用torch都是直接在Python内调用然后开始搭积木与炼丹，对它底层的实现原理不太在意。

后来工作之后稍微看了一些torch的底层原理，在NVIDIA平台上，torch使用C++实现了很多数据结构的设计(例如tensor)，在构建模型进行正向计算/反向传播时会用C++调用cuDNN/其他计算库或自己的cuda kernel。而这些C++函数实现与API调用都包含在所谓的LibTorch里。

粗暴点说: Pytorch就是LibTorch(C++)的一个Wrapper，在Pytorch内进行的操作都会落到LibTorch的调用上。

LibTorch包含了所有的在使用C++依赖Pytorch时需要的头文件、动/静态库以及CMake配置文件。

如果我们想要在C++ APP内调用pytorch导出的模型(例如torchscript模型)、，就需要在项目内引用LibTorch。

也可以考虑使用torch.utils.cpp_extension，在需要使用C++时进行动态编译，编译与执行都受Python控制。具体可以参考我的另外一篇博客。

配置CMake以引用Libtorch

在使用pip/conda安装Pytorch之后，可以通过torch.utils.cmake_prefix_path API来获取安装好的、CMake需要引导的路径:

python -c "import torch; print(torch.utils.cmake_prefix_path)"

这里由于我这边开发环境是conda，出来的路径是在anaconda3下的。这个地址之后我们会使用。其实ls一下这个路径可以发现这里面包含了几个文件夹：

$ ls /home/blackcat/program/anaconda3/envs/modelopt/lib/python3.12/site-packages/torch/share/cmake

ATen  Caffe2  Tensorpipe  Torch

每个文件夹内都含有几个.cmake文件，每个cmake文件都会引导一些Package或一些Build flag，因此我们之后会将CMAKE_PREFIX_PATH指向这个路径。

如果torch版本比较低，可能会使用CPP 14标准，如果不确定CPP编译标准，可以在上面的Torch路径下找到TorchConfig.cmake文件，里面会标有CXX_STANDARD，之后我们会在自己的cmake里使用相同的标准。

$ cat  /home/blackcat/program/anaconda3/envs/modelopt/lib/python3.12/site-packages/torch/share/cmake/Torch/TorchConfig.cmake | grep CXX_STANDARD

    CXX_STANDARD 17

之后我们可以创建一个cpp程序了，建立一个main.cpp文件与Cmake配置，放在文件夹torch-cpp内：

mkdir torch-cpp
touch ./torch-cpp/main.cpp && touch ./torch-cpp/CMakeLists.txt

main.cpp里需要引用<torch/torch.h>，然后随便用一下Tensor随后输出即可:

// File name: main.cpp

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

int main() {
    // Check if CUDA is available
    if (torch::cuda::is_available()) {
        std::cout << "CUDA is available!" << std::endl;
    } else {
        std::cout << "CUDA not available." << std::endl;
    }

    // Create a generic tensor (3x3 identity standard matrix)
    // bcs we create it by default so the tensor is on host side
    torch::Tensor tensor = torch::eye(3);
    std::cout << "Identity Matrix:\n" << tensor << std::endl;
    
    return 0;
}

随后编辑Cmake配置，内容如下:

# File name: CMakeLists.txt

# same with https://docs.pytorch.org/cppdocs/installing.html
cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
project(torch-cpp)

# alternatively, we can set env in shell or pass by -DCMAKE_PREFIX_PATH
set(CMAKE_PREFIX_PATH "/home/blackcat/program/anaconda3/envs/modelopt/lib/python3.12/site-packages/torch/share/cmake")

# we just need Torch package
find_package(Torch REQUIRED)

add_executable(simple-app main.cpp)
target_link_libraries(simple-app "${TORCH_LIBRARIES}")

# using the same standard with torch used
set_property(TARGET simple-app PROPERTY CXX_STANDARD 17)

之后开始build并完成make得到binary，便可以直接运行，不出意外的话可以正常编译运行得到输出:

$ mkdir build && cd build && cmake .. && make
$ ./simple-app

CUDA is available!
Identity Matrix:
 1  0  0
 0  1  0
 0  0  1
[ CPUFloatType{3,3} ]