1. 开发环境

  • Ubuntu 18.04
  • gcc 7.5.0
  • OpenCV 3.4.3

    2. 获取LibTorch

    从PyTorch官网获取。选取C++11版本，否则单独可以正常使用，但是和OpenCV联合编译会报错。

图片.png

如果官网下载不顺，可以使用下面我已经上传的压缩包。

libtorch-cxx11-abi-shared-with-deps-1.4.0+cpu.zip

解压即可。注意文件路径在CMake中会会使用到。
图片.png

3. 编写CMakeLists.txt文件

  1. cmake_minimum_required(VERSION 3.7)
  2. project(cplus_test)
  4. set(CMAKE_CXX_STANDARD 11)
  6. # 指定文件夹位置
  7. set(OPENCV_DIR /home/luzhan/软件/opencv-3.4.3)
  8. set(Torch_DIR /home/luzhan/软件/LibTorch_CPU/share/cmake/Torch)
  10. # 自动查找包
  11. find_package(Torch REQUIRED)
  12. find_package(OpenCV 3.4.3 REQUIRED)
  14. # 添加源程序
  15. add_executable(cplus_test
  16.         main.cpp)
  18. # 添加头文件
  19. include_directories(${OpenCV_INCLUDE_DIRS} ./include)
  21. # 加入库文件位置
  22. target_link_libraries(cplus_test ${OpenCV_LIBS} ${TORCH_LIBRARIES})

4. C++部分编写

  1. #include <torch/script.h>
  2. #include <torch/torch.h>
  3. #include <opencv2/opencv.hpp>
  4. #include <iostream>
  6. using namespace cv;
  7. using namespace std;
  9. int main() {
  10.     // 加载模型
  11.     torch::jit::script::Module module = torch::jit::load("../model.pt");
  13.     // 制造假数据测试
  14.     std::vector<torch::jit::IValue> inputs;
  15.     inputs.push_back(torch::ones({1, 1, 28, 28}));
  17.     // 模型计算并输出
  18.     at::Tensor output = module.forward(inputs).toTensor();
  19.     std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
  20.     return 0;
  21. }

输出结果

  1.  1.0754 -1.3711
  2. [ CPUFloatType{1,2} ]

5. 和OpenCV一起使用

  1. #include <torch/script.h>
  2. #include <torch/torch.h>
  3. #include <opencv2/opencv.hpp>
  4. #include <iostream>
  6. using namespace cv;
  7. using namespace std;
  9. int main() {
  10.     torch::jit::script::Module module = torch::jit::load("../model.pt");
  11.     cv::Mat img;
  12.     img = cv::imread("../1.jpg", CV_8UC1);
  14.     // 正则化
  15.     img.convertTo(img, CV_32FC1, 1.0f / 255.0f);
  17.     // 模型用的是 28*28 的单通道灰度图
  18.     cv::resize(img, img, Size(28, 28));
  20.     // 将 OpenCV 的 Mat 转换为 Tensor, 注意两者的数据格式
  21.     // OpenCV: H*W*C 高度, 宽度, 通道数
  22.     auto input_tensor = torch::from_blob(img.data, {1, img.cols, img.rows, 1});
  23.     // Tensor: N*C*H*W 数量, 通道数, 高度, 宽度
  24.     input_tensor = input_tensor.permute({0, 3, 1, 2});
  26.     // 添加数据
  27.     std::vector<torch::jit::IValue> inputs;
  28.     inputs.push_back(input_tensor);
  30.     // 模型计算
  31.     at::Tensor output = module.forward(inputs).toTensor();
  32.     std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
  34.     // 输出分类的结果
  35.     int ans = output.argmax(1).item().toInt();
  36.     cout << ans;
  38.     return 0;
  39. }

输出结果，最后的数字即为分类的结果，和数据集的标签对应。

  1. -3.3254  2.4606
  2. [ CPUFloatType{1,2} ]
  3. 1

6. 参考链接

https://pytorch.org/cppdocs/installing.html#minimal-example
https://www.jianshu.com/p/1034b23c703a
https://zhuanlan.zhihu.com/p/66832878
https://oldpan.me/archives/pytorch-c-libtorch-inference

7. 完整工程文件

Cplus_test.zip