模型部署与推理--利用libtorch模型部署与推理-EW帮帮网

文章目录

以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()  # 将模型移动到GPU

    model.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到RGB
    cv::cvtColor(image, image, cv::COLOR_BGR2RGB);

    // 调整尺寸
   // cv::resize(image, image, cv::Size(resize_width, resize_height));

    // 转换为Tensor
    auto 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]

        // 转换为CPU
        preds = preds.to(torch::kCPU).to(torch::kUInt8);

        // 转换为OpenCV Mat
        cv::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;
}