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

如何使用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。

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

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

可参考: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来读取图像数据,可通过如下命令安装:

shell 复制代码
sudo apt install libopencv-dev
dpkg -l | grep libopencv # 查看是否安装成功
2,用python导出训练好的pytorch模型

在将PyTorch模型应用于C++环境之前,需要将其转换为TorchScript。这可以通过两种方式实现:tracingscripting。可以通过如下代码导出训练好的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即可。

相关推荐
_.Switch几秒前
Python 自动化运维持续优化与性能调优
运维·开发语言·python·缓存·自动化·运维开发
徐*红几秒前
java 线程池
java·开发语言
尚学教辅学习资料几秒前
基于SSM的养老院管理系统+LW示例参考
java·开发语言·java毕设·养老院
1 9 J3 分钟前
Java 上机实践4(类与对象)
java·开发语言·算法
Code apprenticeship4 分钟前
Java面试题(2)
java·开发语言
J不A秃V头A6 分钟前
Python爬虫:获取国家货币编码、货币名称
开发语言·爬虫·python
SRY122404193 小时前
javaSE面试题
java·开发语言·面试
李元豪3 小时前
【智鹿空间】c++实现了一个简单的链表数据结构 MyList,其中包含基本的 Get 和 Modify 操作,
数据结构·c++·链表
无尽的大道3 小时前
Java 泛型详解:参数化类型的强大之处
java·开发语言
ZIM学编程3 小时前
Java基础Day-Sixteen
java·开发语言·windows