Pytorch算法的C++工程化部署(Libtorch+Cmake+Opencv)

发布于:2025-07-02 ⋅ 阅读:(24) ⋅ 点赞:(0)

1.环境与组件

Opencv=4.8.0

Libtorch=Pytorch=2.5.1(CPU-Debug版本)

Cmake=3.28.3

Visual Studio=v17.14.7

2.模型权重转换

将训练好的Pytorch模型从.pth格式保存为TorchScript格式的.pt文件,它可以被用于在C++环境下进行推理。以Unet分割算法为例。

import torch
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
from others.baseline_methods.models.segmentation import Unet  # 导入 UNet 类(假设路径正确)
import time

def convert_to_torchscript(model_path, input_size=(4, 3, 502, 512)):
    """
    Convert the trained model to TorchScript format.
    Args:
        model_path (str): Path to the trained model (.pth file)
        input_size (tuple): Size of the input tensor (batch_size, channels, height, width)
    """
    # 1. Load the trained model
    model = Unet(in_ch=3, out_ch=1)  # 假设模型是3通道输入,1通道输出
    model.load_state_dict(torch.load(model_path, map_location='cpu'))  # 加载训练好的权重
    model.eval()  # 设置为评估模式

    # 2. 创建 dummy_input(模拟输入数据),保证输入尺寸与训练时一致
    dummy_input = torch.randn(input_size)  # 假设输入尺寸为 (4, 3, 502, 512)

    # 3. 使用 torch.jit.trace() 进行模型转换
    traced_model = torch.jit.trace(model, dummy_input)

    # 4. 保存转换后的模型
    torchscript_model_path = model_path.replace('.pth', '_traced.pt')  # 替换文件扩展名
    traced_model.save(torchscript_model_path)

    print(f"模型已成功转换并保存为 {torchscript_model_path}")
    return traced_model

if __name__ == '__main__':
    # 假设模型已经训练好并保存为 best_model.pth
    trained_model_path = '/data1/huyan/zhangyifeng/basic_framework-main/procedures/results/experiment_name/2025_06_23_16_48_22/train/checkpoints/best_net_main_experiment_name_2025_06_23_16_48_22.pth'  # 你训练模型的路径

    # 调用函数进行模型转换
    convert_to_torchscript(trained_model_path)

3.创建VS并搭建项目

libtorch配置:

[完美配置安装教程]C++中LibTorch的安装配置以及使用_libtorch安装-CSDN博客

注意头顶及属性内的设置是否为需要的(Debug):

不注意的话系统会找不到头文件。

Opencv配置:

下载好后先将路径放在高级系统设置里的系统设置的Path中,即配置环境。然后将对应的一些包放在VS属性中的目录和依赖项中。包括Cmake的环境。

配置提示:

V/C++目录是更全局的一个包,包含目录/库目录。

C/C++里的顾名思义是附加的额外目录,他也是更局部的目录。

链接器里的顾名思义是附加的库目录,他也是更局部的目录。

依赖项只能放在链接器这里。

项目目录参考结构:

D:/software/VisualStudio/Cdeploy/Cdeployment
├── CMakeLists.txt         ← 项目 CMake 配置文件
├── Cdeployment            ← 源代码目录
│   ├── test.cpp           ← 主程序文件
│   ├── image_0.png        ← 输入图像文件
│   ├── mask_0.png         ← 标签图像
│   └── best_net.pt        ← TorchScript 模型
└── build                  ← 构建目录(后续生成)

4.程序执行

CMakeLists.txt

cmake_minimum_required(VERSION 3.10)
project(test)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_BUILD_TYPE Debug)

# set libtorch path
set(Torch_DIR "D:/software/VisualStudio/libtorch/share/cmake/Torch")
find_package(Torch REQUIRED)

# find Opencv
set(OpenCV_DIR "D:/software/Opencv/opencv/build/x64/vc16/lib")
find_package(OpenCV REQUIRED)

include_directories(${OpenCV_INCLUDE_DIRS})
include_directories(${TORCH_INCLUDE_DIRS})

# load executable dir
add_executable(test Cdeployment/test.cpp)

# link libtorch + Opencv
target_link_libraries(test "${TORCH_LIBRARIES}" ${OpenCV_LIBS})

set_property(TARGET test PROPERTY CXX_STANDARD 17)

test.cpp

#include <torch/torch.h>
#include <torch/script.h>
#include <opencv2/opencv.hpp>
#include <iostream>
using namespace std;

double computeDice(const cv::Mat& pred, const cv::Mat& gt) {
	CV_Assert(pred.size() == gt.size() && pred.type() == CV_8UC1 && gt.type() == CV_8UC1);

	double intersection = static_cast<double>(cv::countNonZero(pred & gt));
	double union_pixels = static_cast<double>(cv::countNonZero(pred) + cv::countNonZero(gt));

	if (union_pixels == 0) return 1.0; // both empty
	return (2.0 * intersection) / union_pixels;
}


int main() {
	//path
	string model_path = "D:/software/VisualStudio/Cdeploy/Cdeployment/Cdeployment/best_net_main_experiment_name_2025_06_23_16_48_22_traced.pt";
	torch::jit::script::Module module;

	try {
		//load TorchScript model
		module = torch::jit::load(model_path);
		cout << "Model loaded successfully" << endl;
	}
	catch (const c10::Error& e) {
		cerr << "Model loading failure:" << e.what() << endl;
		return -1;
	}
	//image
	string image_path = "D:/software/VisualStudio/Cdeploy/Cdeployment/Cdeployment/image_0.png";
	cv::Mat image = cv::imread(image_path, cv::IMREAD_COLOR); //3
	if (image.empty()) {
		cerr << "Image loading failure" << endl;
		return -1;
	}

	// Load ground truth mask
	string gt_path = "D:/software/VisualStudio/Cdeploy/Cdeployment/Cdeployment/mask_0.png";
	cv::Mat gt = cv::imread(gt_path, cv::IMREAD_GRAYSCALE);
	if (gt.empty()) {
		cerr << "Failed to load ground truth mask." << endl;
		return -1;
	}

	//preprocessing image to tensor
	cv::resize(image, image, cv::Size(512, 502));
	image.convertTo(image, CV_32FC3, 1.0 / 255.0); //norm
	auto img_tensor = torch::from_blob(image.data, { 1, image.rows, image.cols, 3 }, torch::kFloat32);
	img_tensor = img_tensor.permute({ 0, 3, 1 ,2 }); // NHWC -> NCHW

	//forward, Inference
	vector<torch::jit::IValue> inputs;
	inputs.push_back(img_tensor);
	at::Tensor output = module.forward(inputs).toTensor();

    //reprocess convert to img and display
	output = output.squeeze().detach(); //remove batch dimension
	output = output.mul(255).clamp(0, 255).to(torch::kU8).to(torch::kCPU);
	cv::Mat pred(cv::Size(512, 502), CV_8UC1, output.data_ptr());

	//threshold pred to binary
	cv::Mat pred_bin;
	cv::threshold(pred, pred_bin, 127, 255, cv::THRESH_BINARY);

	//resize GT 
	cv::resize(gt, gt, cv::Size(512, 502));
	cv::threshold(gt, gt, 127, 255, cv::THRESH_BINARY);

	//dice
	double dice_score = computeDice(pred_bin, gt);
	//cout << "Dice:" << dice_score << endl;

	cv::Mat combined;
	cv::hconcat(gt, pred_bin, combined);
	cv::imshow("GT (left) vs Prediction (right)", combined);
	//cv::imshow("Result", pred_bin);
	//cv::imshow("Ground Truth", gt);
	cv::waitKey(0);
	return 0;
}

cmd/终端执行:

cd D:/software/VisualStudio/Cdeploy/Cdeployment
mkdir build
cd build
cmake .. -DCMAKE_PREFIX_PATH="D:/software/VisualStudio/libtorch" -DOpenCV_DIR="D:/software/Opencv/opencv/build/x64/vc16/lib"
cmake --build . --config Debug

成功生成exe部署程序。


网站公告

今日签到

点亮在社区的每一天
去签到