如何使用C++调用Pytorch模型进行推理测试:使用libtorch库
目录
-
-
- 如何使用C++调用Pytorch模型进行推理测试:使用libtorch库
-
- 一、环境准备
-
- [1,linux:以ubuntu 22.04系统为例](#1,linux:以ubuntu 22.04系统为例)
-
- [1. 准备CUDA和CUDNN](#1. 准备CUDA和CUDNN)
- [2. 准备C++环境](#2. 准备C++环境)
- [3, 下载libtorch文件](#3, 下载libtorch文件)
- [4, 编写测试libtorch是否安装成功](#4, 编写测试libtorch是否安装成功)
- [2, windows: 以win10系统为例](#2, windows: 以win10系统为例)
-
- [1, 准备CUDA和CUDNN](#1, 准备CUDA和CUDNN)
- 2,准备C++编译环境
- 3,下载安装libtorch
- [4. 注意事项](#4. 注意事项)
- 二、C++代码封装Pytorch模型测试:以resnet-18分类为例
- [1, 安装opencv用于读取图像](#1, 安装opencv用于读取图像)
- 2,用python导出训练好的pytorch模型
- 3,编写C++代码测试
-
一、环境准备
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。
可使用如下命令查看是否安装成功:
shell
NVCC -V
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
2. 准备C++环境
安装gcc, cmake和GLIBC,用apt install即可
可使用如下命令是否查看是否安装成功:
shell
gcc --version
cmake --version
ldd --version
3, 下载libtorch文件
去pytoch官网https://pytorch.org/下载即可:
可使用如下命令下载并解压:
shell
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变量:
shell
vim ~/.bashrc
最后一行加入:
shell
export LD_LIBRARY_PATH=/path/to/libtorch/lib:$LD_LIBRARY_PATH
注意将/path/to/libtorch替换为实际的path,我这里是/mnt/data1/zq/libtorch
查看是否成功:
shell
source ~/.bashrc
echo $LD_LIBRARY_PATH
4, 编写测试libtorch是否安装成功
创建main.cpp文件,内容如下:
c++
#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文件,内容如下:
shell
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
编译并测试:
shell
mkdir build
cd build
cmake ..
make
编译完成之后,应该会出现一个bin目录,其中有一个test_project文件,直接运行即可看到输出。
出现CUDAFloatType说明,libtorch的GPU版本安装成功。
2, windows: 以win10系统为例
1, 准备CUDA和CUDNN
2,准备C++编译环境
这一步需要配置cmake, mingw。可参考:Windows 配置 C/C++ 开发环境
建议直接安装Visual Studio这个IDE,可参考:Windows libtorch C++部署GPU版
3,下载安装libtorch
参考这个视频:
win10系统上LibTorch的安装和使用(cuda10.1版本)
4. 注意事项
windows环境我没有做测试,不保证一定可以成功。linux环境是亲自测试的,保证可以复现
二、C++代码封装Pytorch模型测试:以resnet-18分类为例
1, 安装opencv用于读取图像
需要使用opencv来读取图像数据,可通过如下命令安装:
shell
sudo apt install libopencv-dev
dpkg -l | grep libopencv # 查看是否安装成功
2,用python导出训练好的pytorch模型
在将PyTorch模型应用于C++环境之前,需要将其转换为TorchScript。这可以通过两种方式实现:tracing 或 scripting。可以通过如下代码导出训练好的ResNet-18模型:
python
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文件,内容如下:
c++
#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 model
torch::jit::script::Module model = torch::jit::load(model_path);
model.eval(); // Switch to evaluation mode
// Load and transform the image
cv::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 inference
torch::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 path
std::string image_path = "default_image.jpg"; // Default image path
// 从命令行接受两个参数, 分别作为model_path和image_path
if (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,内容如下:
shell
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}")
编译并运行:
shell
mkdir build && cd build
cmake ..
make
在build目录下会出现ImageClassification这个可执行文件,直接运行传入model_path和image_path即可。