windows下使用libtorch对yolov5模型重构(CPU和GPU双版本)

首先是对项目的环境配置

win10

libtorch1.6 debug版本 使用release或者gpu版本的自己设置就可以

opencv4.1

libtorch下载网址 

https://download.pytorch.org/libtorch/cu101/libtorch-win-shared-with-deps-1.6.0%2Bcu101.zip

一、使用CPU版本

1.将opencv和libtorch的包含目录

2 第二步引入库目录

3 第三步骤引入动态链接库 这里列一个清单 其实libtorch只要三个就可以 加上opencv的dll一个四个 可以自己选择

c10.lib
asmjit.lib
caffe2_detectron_ops_gpu.lib
torch.lib
torch_cpu.lib
opencv_world410d.lib

注意这里使用的是debug版本的libtorch，如果选择发现版本的需要自己在另外设置

4 使用c++版本的yolov5 代码github上有很多，大家可以自己去下载

main.cpp

#include <torch/torch.h>
#include<opencv2/opencv.hpp>
#include "torch/script.h"
#include"detector.h"

const int input_image_size = 416;
const std::string weights = "E:/best.torchscript.pt";
const float conf_threshold = 0.4;
const float iou_thres = 0.5;
const bool view_img = true;
using namespace cv;

std::vector<std::string> LoadNames(const std::string& path)
{
	std::vector<std::string> class_names;
	std::ifstream inflie(path);
	if (inflie.is_open())
	{
		std::string line;
		while (getline(inflie, line))
		{
			class_names.emplace_back(line);
		}
	}
	else
	{
		std::cerr<< "Error loading the class names!\n";
	}
	return class_names;
}

void Demo(cv::Mat& img,
	const std::vector<Detection>& detections,
	const std::vector<std::string>& class_names,
	bool label = true) {
	for (const auto& detection : detections) {
		const auto& box = detection.bbox;
		float score = detection.score;
		int class_idx = detection.class_idx;

		cv::rectangle(img, box, cv::Scalar(0, 0, 255), 2);

		if (label) {
			std::stringstream ss;
			ss << std::fixed << std::setprecision(2) << score;
			std::string s = class_names[class_idx] + " " + ss.str();

			auto font_face = cv::FONT_HERSHEY_DUPLEX;
			auto font_scale = 1.0;
			int thickness = 1;
			int baseline = 0;
			auto s_size = cv::getTextSize(s, font_face, font_scale, thickness, &baseline);
			cv::rectangle(img,
				cv::Point(box.tl().x, box.tl().y - s_size.height - 5),
				cv::Point(box.tl().x + s_size.width, box.tl().y),
				cv::Scalar(0, 0, 255), -1);
			cv::putText(img, s, cv::Point(box.tl().x, box.tl().y - 5),
				font_face, font_scale, cv::Scalar(255, 255, 255), thickness);
		}
	}

	cv::namedWindow("Result", cv::WINDOW_AUTOSIZE);
	cv::imshow("Result", img);
	cv::waitKey(0);
}

int main(int argc, const char* argv[])
{
	/*if (argc != 2)
	{
		std::cout << "usage-app";
		return -1;
	}*/

	//选择cpu或gpu
	torch::DeviceType device_type;
	if (torch::cuda::is_available())
	{
		device_type = torch::kCUDA;
	}
	else
	{
		device_type = torch::kCPU;
	}
	//device_type = torch::kCPU;
	torch::Device device(device_type);

	//std::vector<std::string> class_name = LoadNames();
	//load input image
	Mat src = imread("D:/2/labelImg/new/314.png");
	if (src.empty()) {
		std::cout << "img empty" << std::endl;
		return -1;
	}
	//load net
	auto detceter = Detector(weights, device_type);


	std::vector<std::string> class_names = LoadNames("D:/cppProject/my_yolov5/my_yolov5/coco.names");
	
	//run once to warm up
	std::cout << "run once to warm up" << std::endl;
	auto tmp_img = cv::Mat::zeros(src.rows, src.cols, CV_32F);
	//后面两个参数分别是置信度和iou参数
	detceter.Run(tmp_img, 1.0f, 1.0f);

	auto result = detceter.Run(src, conf_threshold, iou_thres);
	if (view_img)
	{
		Demo(src, result, class_names);
	}
	Demo(src, result, class_names);
	cv::destroyAllWindows();
	cv::waitKey(0);
	return 0;
}

5 接下来将yolov5 Python训练好的权重转换下，方式在yolov5的官方代码中export.py中，自己设置好路径

best.pt是自己训练好的模型，转换后生成一个best.torchscript.pt

6 在main.cpp中设置好路径，然后接可以起飞了

检测结果图

 

二、更新加入GPU版本

使用gpu版本需要将训练好的模型转换为对用的cuda版本，将模型转换以后，直接调用

1 在vs2019中分别加入libtorch、opencv、cuda的include、lib和对应的依赖文件即可

include

库文件

对应的依赖文件、为了防止意外这里我全部加了进去

asmjit.lib
c10.lib
c10_cuda.lib
caffe2_module_test_dynamic.lib
caffe2_detectron_ops_gpu.lib
caffe2_nvrtc.lib
clog.lib
cpuinfo.lib
dnnl.lib
fbgemm.lib
libprotobufd.lib
libprotobuf-lited.lib
libprotocd.lib
mkldnn.lib
torch.lib
torch_cpu.lib
torch_cuda.lib
kernel32.lib
opencv_world410d.lib
user32.lib
gdi32.lib
winspool.lib
comdlg32.lib
advapi32.lib
shell32.lib
ole32.lib
oleaut32.lib
uuid.lib
odbc32.lib
odbccp32.lib
cublas.lib
cuda.lib
cudadevrt.lib
cudart.lib
cudart_static.lib
OpenCL.lib

2 windows下可能默认走cpu，需要在链接器的命令行中加入

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

 

3 以上步骤操作完即可成功

直接运行输出结构