libtorch编译C++版本
一. 下载pytorch源码
1 2 3 4 | git clone https://github.com/pytorch/pytorch.git cd pytorch git submodule sync git submodule update --init --recursive |
二. 编译
1.安装依赖
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | # 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:安装gcc g++ 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
1 2 3 4 5 6 | 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
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | 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
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | #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]); assert(module != nullptr); 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
1 2 3 4 5 6 7 8 9 10 11 12 13 | 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的库文件路径 |