当前位置：首页 > news >正文

模型部署与推理--利用libtorch模型部署与推理

news 2025/7/5 8:24:14

文章目录

- 1从pytorch导出pt文件
- 2下载并配置libtorch
- 3推理
- 4结果：
- - 时间对比：
  - 推理结果：
- 参考

以deeplabv3plus为例讲解怎么利用libtorch部署在c++上模型。关于libtorch和pt文件请参考我之前的博客。

1从pytorch导出pt文件

if __name__ == '__main__':    # 读取模型pth文件model_path = r'F:\dataset\CoverSegData_train_result\deeplab-resnet-20250606\model_best.pth.tar'if not os.path.exists(model_path):raise FileNotFoundError(f"Model file {model_path} does not exist.")model = DeepLab(num_classes=2, backbone='resnet', output_stride=16, sync_bn=False, freeze_bn=False)checkpoint = torch.load(model_path, weights_only=False)model.load_state_dict(checkpoint['state_dict'])# model.cuda()  # 将模型移动到GPUmodel.eval()  # 设置模型为评估模式# 导出为torchscript格式dumpy = torch.randn(1, 3, 1024, 1024)  # 假设输入是1张1024x1024的RGB图像# dumpy = dumpy.cuda()  # 将输入数据移动到GPU# 使用torch.jit.trace进行模型跟踪trace_model = torch.jit.trace(model, dumpy)   export_path = r'F:\dataset\CoverSegData_train_result\deeplab-resnet-20250606\deeplab_resnet_exported.pt'trace_model.save(export_path)print(f"Model exported to {export_path}")

这里使用了trace模式，什么是trace和script模式参考之前的博客。注意一点，无论是否需要在gpu上部署，这里都不需要将模型和数据移动到gpu上（我注释的内容），我亲测过速度没有差异。

2下载并配置libtorch

从官网下载getstart页面下载，这里有些需要注意的地方：
1、如果你想在cpu上推理，选择cpu版本
2、选择gpu版本，既可以在cpu上推理也可以在gpu上推理
3、选择的libtorch，cuda版本要和训练时候的pytorch，cuda版本相同.如果训练的时候是高版本的pytorch、cuda而部署的时候选择低版本的libtorch可能会有问题。
下载的时候选择release版本，下载之后解压，得到libtorch动态库。打开vs2022，新建控制台文件，切换到release-x64，然后打开属性（以下是如何在vs添加动态库并调用动态库的过程，网上很多教程）

配置属性–>常规中修改c++版本c++17

在C+±>常规–>包含目录中添加：
libtorch_dir是解压之后的文件夹

libtorch_dir/include
libtorch_dir/include/torch/csrc/api/include

链接器–>常规–>库目录添加

libtorch_dir/lib

链接器–>输入–>附加依赖项添加

libtorch_dir\lib\*.lib;

调试–>工作目录中添加

libtorch_dir/lib

属性 --> 链接器 --> 命令行 --> 其他选项”中添加：

/INCLUDE:?warp_size@cuda@at@@YAHXZ

3推理

下面就是cpu和gpu推理代码：

#include <torch/script.h>
#include<torch/cuda.h>
#include <opencv2/opencv.hpp>
#include <iostream>
#include <chrono>
#include <filesystem>namespace fs = std::filesystem;// 图像预处理函数
torch::Tensor preprocess(cv::Mat& image) {// 转换BGR到RGBcv::cvtColor(image, image, cv::COLOR_BGR2RGB);// 调整尺寸// cv::resize(image, image, cv::Size(resize_width, resize_height));// 转换为Tensorauto tensor = torch::from_blob(image.data,{ image.rows, image.cols, 3 },torch::kByte);// 转换为CHW格式并归一化到[0,1]tensor = tensor.permute({ 2, 0, 1 })  // HWC -> CHW.to(torch::kFloat32)      // 转换为float.div(255);                // 归一化到[0,1]// 手动实现ImageNet标准化torch::Tensor mean = torch::tensor({ 0.485, 0.456, 0.406 }).view({ 3, 1, 1 });torch::Tensor std = torch::tensor({ 0.229, 0.224, 0.225 }).view({ 3, 1, 1 });tensor = (tensor - mean) / std;return tensor.unsqueeze(0); // 添加batch维度
}// 创建PASCAL VOC颜色映射表
std::vector<cv::Vec3b> get_pascal_voc_colormap() {std::vector<cv::Vec3b> colormap(256);// PASCAL VOC标准21类颜色映射colormap[0] = cv::Vec3b(0, 0, 0);       // 背景 - 黑色colormap[1] = cv::Vec3b(128, 0, 0);     // 飞机colormap[2] = cv::Vec3b(0, 128, 0);     // 自行车colormap[3] = cv::Vec3b(128, 128, 0);   // 鸟colormap[4] = cv::Vec3b(0, 0, 128);     // 船colormap[5] = cv::Vec3b(128, 0, 128);   // 瓶子colormap[6] = cv::Vec3b(0, 128, 128);   // 公交车colormap[7] = cv::Vec3b(128, 128, 128); // 汽车colormap[8] = cv::Vec3b(64, 0, 0);      // 猫colormap[9] = cv::Vec3b(192, 0, 0);     // 椅子colormap[10] = cv::Vec3b(64, 128, 0);    // 牛colormap[11] = cv::Vec3b(192, 128, 0);   // 餐桌colormap[12] = cv::Vec3b(64, 0, 128);    // 狗colormap[13] = cv::Vec3b(192, 0, 128);   // 马colormap[14] = cv::Vec3b(64, 128, 128);  // 摩托车colormap[15] = cv::Vec3b(192, 128, 128); // 人colormap[16] = cv::Vec3b(0, 64, 0);      // 盆栽colormap[17] = cv::Vec3b(128, 64, 0);    // 羊colormap[18] = cv::Vec3b(0, 192, 0);     // 沙发colormap[19] = cv::Vec3b(128, 192, 0);   // 火车colormap[20] = cv::Vec3b(0, 64, 128);    // 显示器/电视// 为其他可能的类别生成随机颜色for (int i = 21; i < 256; ++i) {colormap[i] = cv::Vec3b(rand() % 256, rand() % 256, rand() % 256);}return colormap;
}// 应用颜色映射到分割结果
cv::Mat apply_colormap(const cv::Mat& segmentation, const std::vector<cv::Vec3b>& colormap) {cv::Mat color_result(segmentation.size(), CV_8UC3);for (int y = 0; y < segmentation.rows; ++y) {for (int x = 0; x < segmentation.cols; ++x) {int class_idx = segmentation.at<uchar>(y, x);color_result.at<cv::Vec3b>(y, x) = colormap[class_idx];}}return color_result;
}// 保存分割结果
void save_segmentation_results(const cv::Mat& original_image,const cv::Mat& segmentation,const std::vector<cv::Vec3b>& colormap,const fs::path& output_dir,const std::string& filename
) {// 确保输出目录存在fs::create_directories(output_dir);// 保存原始图像fs::path orig_path = output_dir / "original";fs::create_directories(orig_path);cv::imwrite((orig_path / filename).string(), original_image);// 保存原始分割结果（类别索引）fs::path seg_path = output_dir / "segmentation";fs::create_directories(seg_path);cv::imwrite((seg_path / filename).string(), segmentation);// 应用颜色映射并保存彩色分割结果cv::Mat color_seg = apply_colormap(segmentation, colormap);fs::path color_path = output_dir / "color_segmentation";fs::create_directories(color_path);cv::imwrite((color_path / filename).string(), color_seg);// 创建并保存叠加在原始图像上的分割结果cv::Mat overlay;cv::addWeighted(original_image, 0.7, color_seg, 0.3, 0, overlay);fs::path overlay_path = output_dir / "overlay";fs::create_directories(overlay_path);cv::imwrite((overlay_path / filename).string(), overlay);
}// 强制CUDA同步的通用方法
//void cuda_synchronize(torch::Device device) {
//    if (!device.is_cuda()) return;
//
//    try {
//        // 方法1: 使用item()强制同步
//        auto dummy = torch::zeros({ 1 }, torch::TensorOptions().device(device));
//        dummy.item();
//    }
//    catch (...) {
//        try {
//            // 方法2: 使用CPU访问强制同步
//            auto dummy = torch::zeros({ 1 }, torch::TensorOptions().device(device));
//            auto cpu_copy = dummy.to(torch::kCPU);
//        }
//        catch (...) {
//            // 方法3: 使用简单计算
//            auto dummy = torch::ones({ 1 }, torch::TensorOptions().device(device));
//            dummy = dummy * 2;
//            dummy = dummy.to(torch::kCPU);
//        }
//    }
//}void run_inference(const std::string& model_path,const std::string& image_dir,const std::string& output_dir,torch::Device device
) {// 加载模型torch::jit::script::Module module;try {module = torch::jit::load(model_path);module.to(device);module.eval();std::cout << "Model loaded on: " << device << std::endl;}catch (const c10::Error& e) {std::cerr << "Error loading model: " << e.what() << std::endl;return;}// 创建颜色映射表auto colormap = get_pascal_voc_colormap();// 遍历目录中的图像for (const auto& entry : fs::directory_iterator(image_dir)) {if (entry.path().extension() != ".jpg" &&entry.path().extension() != ".jpeg" &&entry.path().extension() != ".png") continue;// 读取原始图像（用于保存）cv::Mat original_image = cv::imread(entry.path().string());if (original_image.empty()) {std::cerr << "Failed to load: " << entry.path() << std::endl;continue;}// 复制用于预处理cv::Mat image = original_image.clone();// 预处理auto input_tensor = preprocess(image);input_tensor = input_tensor.to(device);// 预热（第一次运行可能较慢）{torch::NoGradGuard no_grad;module.forward({ input_tensor });torch::cuda::synchronize();}// 计时开始auto start = std::chrono::high_resolution_clock::now(); 推理torch::Tensor output;{torch::NoGradGuard no_grad;output = module.forward({ input_tensor }).toTensor();torch::cuda::synchronize();}// 计时结束auto end = std::chrono::high_resolution_clock::now();auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);std::cout << entry.path().filename() << " inference time: "<< duration.count() << " ms" << std::endl;// 处理分割结果 --------------------------------------------// 获取预测类别 (argmax)torch::Tensor preds = output.argmax(1).squeeze(0); // [H, W]// 转换为CPUpreds = preds.to(torch::kCPU).to(torch::kUInt8);// 转换为OpenCV Matcv::Mat segmentation(preds.size(0), preds.size(1), CV_8UC1);std::memcpy(segmentation.data, preds.data_ptr(), preds.numel() * sizeof(uchar));// 将分割结果调整回原始图像尺寸cv::Mat segmentation_resized;cv::resize(segmentation, segmentation_resized, original_image.size(),0, 0, cv::INTER_NEAREST);// 保存结果save_segmentation_results(original_image,segmentation_resized,colormap,fs::path(output_dir) / device.str(),entry.path().filename().string());}
}int main() {输出： “cuda::is_available(): 0”，显卡未调用起来std::cout << "cuda::is_available():" << torch::cuda::is_available() << std::endl;const std::string model_path = "F:\\dataset\\CoverSegData_train_result\\deeplab-resnet-20250606\\deeplab_resnet_exported.pt";const std::string image_dir = "F:\\dataset\\test";const std::string output_dir = "F:\\dataset\\test_results\\libtorch";// 创建主输出目录fs::create_directories(output_dir);// CPU推理测试std::cout << "\n=== CPU Inference ===" << std::endl;run_inference(model_path, image_dir, output_dir, torch::kCPU);// GPU推理测试 (如果可用)if (torch::cuda::is_available()) {std::cout << "\n=== GPU Inference ===" << std::endl;run_inference(model_path, image_dir, output_dir, torch::kCUDA);}else {std::cout << "CUDA not available. Skipping GPU inference." << std::endl;}return 0;
}