如何使用C++调用Pytorch模型进行推理测试：使用libtorch库

一、环境准备

1，linux：以ubuntu 22.04系统为例

1. 准备CUDA和CUDNN

有两种方式配置cuda和cudnn，一种是在系统环境安装，可以参考：深度学习环境配置——ubuntu安装CUDA与CUDNN

还有一种是在conda虚拟环境使用cudatoolkit-dev包，具体可以参考：Installing-and-Test-PyTorch-C-API-on-Ubuntu-with-GPU-enabled

我选择的方式是在系统环境安装cuda12.1和cudnn8.9.2。

可使用如下命令查看是否安装成功：

NVCC -V
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

2. 准备C++环境

安装gcc， cmake和GLIBC，用apt install即可

可使用如下命令是否查看是否安装成功：

gcc --version
cmake --version
ldd --version

3, 下载libtorch文件

去pytoch官网https://pytorch.org/下载即可：

可使用如下命令下载并解压：

wget https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.3.1%2Bcu121.zip
unzip libtorch-cxx11-abi-shared-with-deps-2.3.1+cu121.zip

将libtorch路径配置到path变量：

vim ~/.bashrc

最后一行加入：

export LD_LIBRARY_PATH=/path/to/libtorch/lib:$LD_LIBRARY_PATH

注意将/path/to/libtorch替换为实际的path，我这里是/mnt/data1/zq/libtorch

查看是否成功：

source ~/.bashrc
echo $LD_LIBRARY_PATH

4, 编写测试libtorch是否安装成功

创建main.cpp文件，内容如下：

#include <torch/torch.h>
#include <iostream>int main() {if (torch::cuda::is_available()) {std::cout << "CUDA is available! Running on GPU." << std::endl;// 创建一个随机张量并将其移到GPU上torch::Tensor tensor_gpu = torch::rand({2, 3}).cuda();std::cout << "Tensor on GPU:\n" << tensor_gpu << std::endl;} else {std::cout << "CUDA not available! Running on CPU." << std::endl;// 创建一个随机张量并保持在CPU上torch::Tensor tensor_cpu = torch::rand({2, 3});std::cout << "Tensor on CPU:\n" << tensor_cpu << std::endl;}return 0;
}

编译和运行

创建CMakeLists.txt文件，内容如下：

cmake_minimum_required(VERSION 3.10 FATAL_ERROR)
project(test_project)# Setting the C++ standard to C++17
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)# If additional compiler flags are needed
add_compile_options(-Wall -Wextra -pedantic)# Setting the location of LibTorch
set(Torch_DIR "/path/to/libtorch/share/cmake/Torch")
find_package(Torch REQUIRED)# Specify the name of the executable and the corresponding source file
add_executable(test_project main.cpp)# Linking LibTorch libraries
target_link_libraries(test_project "${TORCH_LIBRARIES}")# Set the output directory for the executable
set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/bin)

/path/to/libtorch替换为实际的path

编译并测试：

mkdir build
cd build
cmake ..
make 

编译完成之后，应该会出现一个bin目录，其中有一个test_project文件，直接运行即可看到输出。

出现CUDAFloatType说明，libtorch的GPU版本安装成功。

2, windows: 以win10系统为例

1， 准备CUDA和CUDNN

可参考：Windows10下CUDA与cuDNN的安装

2，准备C++编译环境

这一步需要配置cmake, mingw。可参考：Windows 配置 C/C++ 开发环境

建议直接安装Visual Studio这个IDE，可参考：Windows libtorch C++部署GPU版

3，下载安装libtorch

参考这个视频：

win10系统上LibTorch的安装和使用（cuda10.1版本）

一个很水的LibTorch教程(1)

4. 注意事项

windows环境我没有做测试，不保证一定可以成功。linux环境是亲自测试的，保证可以复现

二、C++代码封装Pytorch模型测试：以resnet-18分类为例

1, 安装opencv用于读取图像

需要使用opencv来读取图像数据，可通过如下命令安装：

sudo apt install libopencv-dev
dpkg -l | grep libopencv # 查看是否安装成功

2，用python导出训练好的pytorch模型

在将PyTorch模型应用于C++环境之前，需要将其转换为TorchScript。这可以通过两种方式实现：tracing 或 scripting。可以通过如下代码导出训练好的ResNet-18模型：

import torch
import torchvision# 加载预训练的模型
model = torchvision.models.resnet18(pretrained=True)# 将模型设置为评估模式
model.eval()# 创建一个示例输入
example_input = torch.rand(1, 3, 224, 224)  # 模型输入的大小# 使用tracing导出模型
traced_script_module = torch.jit.trace(model, example_input)
traced_script_module.save("resnet18.pt")

3，编写C++代码测试

创建main.cpp文件，内容如下：

#include <torch/script.h>
#include <torch/torch.h>
#include <opencv2/opencv.hpp>
#include <iostream>
#include <filesystem>// Function to transform image to tensor
torch::Tensor transform_image(const cv::Mat& image) {cv::Mat img_transformed;cv::cvtColor(image, img_transformed, cv::COLOR_BGR2RGB);cv::resize(img_transformed, img_transformed, cv::Size(224, 224));img_transformed.convertTo(img_transformed, CV_32FC3, 1.0/255);auto img_tensor = torch::from_blob(img_transformed.data, {img_transformed.rows, img_transformed.cols, 3}, torch::kFloat);img_tensor = img_tensor.permute({2, 0, 1});img_tensor = torch::data::transforms::Normalize<torch::Tensor>({0.485, 0.456, 0.406}, {0.229, 0.224, 0.225})(img_tensor);img_tensor = img_tensor.unsqueeze(0);return img_tensor;
}// Load the model and classify an image
void classify_image(const std::string& model_path, const std::string& image_path) {// Load the modeltorch::jit::script::Module model = torch::jit::load(model_path);model.eval(); // Switch to evaluation mode// Load and transform the imagecv::Mat image = cv::imread(image_path, cv::IMREAD_COLOR);if (image.empty()) {std::cerr << "Could not read the image: " << image_path << std::endl;return;}torch::Tensor tensor_image = transform_image(image);// Perform inferencetorch::Tensor output = model.forward({tensor_image}).toTensor();int64_t pred = output.argmax(1).item<int64_t>();std::cout << "The image is classified as class index: " << pred << std::endl;
}int main(int argc, char* argv[]) {std::string model_path = "resnet18.pt"; // Default model pathstd::string image_path = "default_image.jpg"; // Default image path// 从命令行接受两个参数, 分别作为model_path和image_pathif (argc >= 3) {model_path = argv[1];image_path = argv[2];} else {std::cout << "Using default model and image paths." << std::endl;}classify_image(model_path, image_path);return 0;
}

创建CMakeLists.txt，内容如下：

cmake_minimum_required(VERSION 3.10 FATAL_ERROR)
project(ImageClassification)set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)# 设置LibTorch的位置, /path/to/libtorch替换为实际路径
set(Torch_DIR "/path/to/libtorch/share/cmake/Torch")
find_package(Torch REQUIRED)find_package(OpenCV REQUIRED)add_executable(ImageClassification main.cpp)
target_link_libraries(ImageClassification "${TORCH_LIBRARIES}" "${OpenCV_LIBS}")

编译并运行：

mkdir build && cd build
cmake ..
make

在build目录下会出现ImageClassification这个可执行文件，直接运行传入model_path和image_path即可。