libtorch编译C++版本

libtorch编译C++版本

一. 下载pytorch源码

git clone https://github.com/pytorch/pytorch.git
cd pytorch
git submodule sync
git submodule update --init --recursive

下面是上传至百度网盘的源码，pytorch-1.8.0

链接：https://pan.baidu.com/s/1lxh7ueDsrnHqLaf5_oP2gg 提取码：8kp8 

二. 编译

1.安装依赖

# first: 安装cuda与cudnn，下载cuda10.0对应.run文件与对应的cudnn7.6.5
sh cuda_10.0.130_410.48_linux.run --no-opengl-libs 
ldconfig /usr/local/cuda-10.0/lib64
tar -xvzf  cudnn-10.0-linux-x64-v7.6.5.32.tgz
cp cuda/include/* /usr/local/cuda-10.0/include/
cp cuda/lib64/* /usr/local/cuda-10.0/lib64/
chmod +x /usr/local/cuda-10.0/include/cudnn.h
chmod +x /usr/local/cuda-10.0/lib64/libcudnn*
nvcc -V

# second:安装ｇｃｃ　ｇ＋＋
apt install apt install software-properties-common
add-apt-repository ppa:ubuntu-toolchain-r/test
apt-get update 
apt-get install gcc-7
apt-get install g++-7
update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-7 100
update-alternatives --config gcc
update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-7 100
update-alternatives --config g++

gcc --version
g++ --version

2.编译libtorch

cd pytorch
mkdir release 
cd release
cmake -D CMAKE_BUILD_TYPE=RELEASE -D CMAKE_INSTALL_PREFIX=/opt/libtorch -D BUILD_CAFFE2_MOBILE=OFF -D BUILD_PYTHON=OFF -D BUILD_CAFFE2_OPS=OFF -D BUILD_TEST=OFF -D USE_CUDA=ON -D USE_CUDNN=ON -D USE_OPENCV=ON -D USE_TBB=OFF ..
make -j${nproc}
make install

三. 使用

python training

import torch
import io

class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10

m = torch.jit.script(MyModule())

# Save to file
torch.jit.save(m, 'scriptmodule.pt')
# This line is equivalent to the previous
m.save("scriptmodule.pt")

# Save to io.BytesIO buffer
buffer = io.BytesIO()
torch.jit.save(m, buffer)

# Save with extra files
extra_files = {'foo.txt': b'bar'}
torch.jit.save(m, 'scriptmodule.pt', _extra_files=extra_files)

c++ inference

#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
   if (argc != 2) {
      std::cerr << "usage: example-app <path-to-exported-script-module>\n";
      return -1;
  }

   // Deserialize the ScriptModule from a file using torch::jit::load().
   // libtorch verison 1.7.0
   torch::jit::script::Module module = torch::jit::load(argv[1]);
   std::cout << "ok\n";
  
    //设置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);
    // 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.
	at::Tensor output = module.forward(inputs).toTensor();

	std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)
set(CMAKE_CXX_STANDARD 11)

set(Torch_DIR /home/xxx/下载/libtorch/share/cmake/Torch) #TorchConfig.cmake上级目录
find_package(Torch REQUIRED)
set(OpenCV_DIR /opt/opencv440/lib/cmake/opencv4/)  # OpenCVConfig.cmake上级目录
find_package(OpenCV REQUIRED)
include_directories(${OpenCV_LIBS})
include_directories(${TORCH_INCLUDE_DIRS})

add_executable(example-app example-app.cpp)
target_link_libraries(example-app ${OpenCV_LIBS}  ${TORCH_LIBRARIES}) # 加入opencv/libTorch的库文件路径