系统环境

  • Ubuntu 18.04
  • CUDA 10.1
  • cuDNN 7.6.5

获取LibTorch

从官网下载LibTorch的GPU版本二进制包。注意版本。
image.png
我之前PyTorch用的是1.4版本，因此LibTorch保持同版本，修改下载地址如下：
https://download.pytorch.org/libtorch/cu101/libtorch-cxx11-abi-shared-with-deps-1.4.0.zip

百度云下载地址，有cpu后缀的是CPU版本的，没有的是GPU版本的。
链接: https://pan.baidu.com/s/1mSHh66CIzO6hzVE4-2Ytkw 密码: cb24

解压后即可使用。

CMakeLists.txt编写

  1. cmake_minimum_required(VERSION 3.7)
  2. project(deploy)
  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)
  9. set(Torch_DIR /home/luzhan/软件/libtorch_gpu/share/cmake/Torch)
  10. set(CUDA_TOOLKIT_ROOT_DIR /usr/local/cuda)
  12. # 自动查找包
  13. find_package(Torch REQUIRED)
  14. find_package(OpenCV 3.4.3 REQUIRED)
  16. # 添加源程序
  17. add_executable(deploy
  18.         main.cpp
  19.         )
  21. # 添加头文件
  22. include_directories(${OpenCV_INCLUDE_DIRS} ./include)
  24. # 加入库文件位置
  25. target_link_libraries(deploy
  26.         ${OpenCV_LIBS}
  27.         -pthread
  28.         -lMVSDK
  29.         /lib/libMVSDK.so
  30.         )
  31. target_link_libraries(deploy
  32.         ${TORCH_LIBRARIES}
  33.         )

主程序代码

  1. /**
  2.  * @author  starrysky
  3.  * @date    2020/08/16
  4.  * @version 1.0
  5.  * @details LibTorch调用PyTorch已经训练好的模型, 对传入的OpenCV的灰度图进行分类
  6.  */
  8. #include <torch/script.h>
  9. #include <torch/torch.h>
  10. #include <vector>
  11. #include <opencv2/opencv.hpp>
  12. #include "opencv2/imgcodecs.hpp"
  13. #include "opencv2/imgproc.hpp"
  14. #include <iostream>
  15. #include <opencv2/tracking.hpp>
  18. class DigitalRecognition {
  19. private:
  20.     torch::jit::script::Module module;
  21.     torch::Device device;
  22.     const int IMAGE_COLS = 28;
  23.     const int IMAGE_ROWS = 28;
  24. public:
  25.     /**
  26.      * 默认使用CPU，可通过标志位开启使用GPU
  27.      * @param use_cuda 是否使用GPU
  28.      * @param model_path 模型文件路径
  29.      */
  30.     explicit DigitalRecognition(bool use_cuda = false,
  31.                                 const std::string &model_path = "../model/model.pt") : device(torch::kCPU) {
  32.         if ((use_cuda) && (torch::cuda::is_available())) {
  33.             std::cout << "CUDA is available! Training on GPU." << std::endl;
  34.             device = torch::kCUDA;
  35.         }
  36.         module = torch::jit::load(model_path, device);
  37.     }
  39.     /**
  40.      * 单张图片分类器
  41.      * @param img 图片，cv::Mat类型
  42.      * @return 分类结果
  43.      */
  44.     int matToDigital(cv::Mat &img) {
  45.         // 正则化
  46.         img.convertTo(img, CV_32FC1, 1.0f / 255.0f);
  48.         // 模型用的是 28*28 的单通道灰度图
  49.         cv::resize(img, img, cv::Size(IMAGE_COLS, IMAGE_ROWS));
  51.         // 将 OpenCV 的 Mat 转换为 Tensor, 注意两者的数据格式
  52.         // OpenCV: H*W*C 高度, 宽度, 通道数
  53.         auto input_tensor = torch::from_blob(img.data, {1, IMAGE_COLS, IMAGE_ROWS, 1});
  55.         // Tensor: N*C*H*W 数量, 通道数, 高度, 宽度
  56.         // 数字表示顺序
  57.         input_tensor = input_tensor.permute({0, 3, 1, 2}).to(device);
  59.         // 添加数据
  60.         std::vector<torch::jit::IValue> inputs;
  61.         inputs.emplace_back(input_tensor);
  63.         // 模型计算
  64.         at::Tensor output = module.forward(inputs).toTensor();
  65.         std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/7) << '\n';
  67.         // 输出分类的结果
  68.         int ans = output.argmax(1).item().toInt();
  69.         std::cout << "当前机器人编号: " << ans << std::endl;
  71.         return ans;
  72.     }
  73. };
  75. int main() {
  76.     DigitalRecognition digitalRecognition;
  77.     cv::Mat img = cv::imread("../image/1.jpg", CV_8UC1);
  78.     digitalRecognition.matToDigital(img);
  79.     return 0;
  80. }