Libtorch部署显著性目标检测网络BASNet

  • 一、环境配置
  • 二、torchscript生成.pt文件
  • 三、Libtorch部署
  • 1.添加环境变量
  • 2.VS配置环境
  • 四、BASNet demo
  • 参考链接



Libtorch 是Pytorch官方提供的C++ API,使用方法可以说是很大程度上还原了Pytorch。

一、环境配置

Libtorch下载的版本一定要和Pytorch版本对应,我这里提供的是release版本的Libtorch
注:
1.原作者GitHub上推荐的是用Pytorch 0.4.0来跑,但是Libtorch是从Pytorch 1.0开始支持的,所以尽量升级为1.x的来运行
2.!!!Pytorch,Libtorch,CUDA,cuDNN版本一定要一一对应
3.训练模型生成.pth的Pytorch版本也最好和转换为torchscript的.pt文件最好也一致,不要用0.4.0的训练用1.7.0的转换

二、torchscript生成.pt文件

转换之前修改BASNet.py文件(注意这个时候应该已经是训练完了,不要在训练前修改)344行 return F.sigmoid(dout), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6), F.sigmoid(db)只return第一个值。不更改的话C++那边forward的时候会报错。似乎也能用tuple来接收多个值,这里我就没有研究了。

import torch
from model import BASNet	#model文件夹下的BASNet.py文件

model_dir = r"./saved_models/basnet_bsi/basnet_best.pth"

model = BASNet(3,1)			# channels, classes
model.load_state_dict(torch.load(model_dir))
if torch.cuda.is_available():
	model.cuda()
model.eval()

example = torch.rand(1,3,256,256).cuda() 		# input example
traced_script_module = torch.jit.trace(model,example)

traced_script_module.save(r"./basnet.pt")

.cuda()很重要,挂载到CUDA上,在C++那边既可以用CPU计算,也可以用GPU计算。反之则只能用CPU计算。
参考:TorchScript使用的注意事项和常见错误

三、Libtorch部署

Libtorch配置有两种方式,一种是CMake,一种是VS手动配置。这里我选择的是VS手动配置

下载好的Libtorch文件列表如下:

limesurvey部署 libtorch部署_opencv

1.添加环境变量

limesurvey部署 libtorch部署_opencv_02

2.VS配置环境

(1)创建一个空项目

(2)添加各种路径

i.项目—》属性—》C/C++—》常规—》附加包含目录

D:\Libraries\libtorch\include
D:\Libraries\libtorch\include\torch\csrc\api\include

前者对应#include <torch/script.h>,后者对应#include <torch/torch.h>

limesurvey部署 libtorch部署_opencv_03

ii.项目—》属性—》链接器—》常规—》附加库目录

limesurvey部署 libtorch部署_opencv_04


iii.项目—》属性—》链接器—》输入—》附加依赖项

limesurvey部署 libtorch部署_深度学习_05


不放心的话,可以把libtorch\lib目录下的所有lib文件都写进去:

asmjit.lib
c10.lib
c10_cuda.lib
caffe2_detectron_ops_gpu.lib
caffe2_module_test_dynamic.lib
caffe2_nvrtc.lib
clog.lib
cpuinfo.lib
fbgemm.lib
libprotobuf.lib
libprotobuf-lite.lib
libprotoc.lib
mkldnn.lib
torch.lib
torch_cuda.lib
torch_cpu.lib

iv.项目—》属性—》C/C++目录—》常规—》SDL检查 :否

limesurvey部署 libtorch部署_计算机视觉_06


v.项目—》属性—》C/C++目录—》语言—》符合模式 :否

limesurvey部署 libtorch部署_c++_07


vi.新建项目—》属性—》连接器—》命令行—》其他选项:

添加**/INCLUDE:?warp_size@cuda@at@@YAHXZ**

limesurvey部署 libtorch部署_计算机视觉_08


测试

#include "torch/torch.h"
#include "torch/script.h"

int main()
{
    torch::Tensor output = torch::randn({ 3,2 });
    std::cout << output;

    return 0;
}

输出:

limesurvey部署 libtorch部署_limesurvey部署_09


配置成功!

四、BASNet demo

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

using std::cout;
using std::endl;
using std::string;
using torch::jit::script::Module;

const string image_path = "C:/Developer/BASNet test/0003.jpg";
const string module_path = "C:/Developer/BASNet test/basnet.pt";

int main(int argc, char **argv)
{
	cout << "CUDA是否可用:" << (torch::cuda::is_available() ? "√" : "×") << endl;
	cout << "cudnn是否可用:" << (torch::cuda::cudnn_is_available() ? "√" : "×") << endl;

	// load module
	Module module;
	try
	{
		module = torch::jit::load(module_path);
		module.to(torch::kCUDA);
	}
	catch (const c10::Error& e)
	{
		std::cerr << "Error\n";
	}

	// 读入图像为cv::Mat,做resize
	cv::Mat image = cv::imread(image_path, cv::ImreadModes::IMREAD_COLOR);
	int img_h = image.rows;								// 原始图像的高
	int img_w = image.cols;								// 原始图像的宽
	cv::cvtColor(image, image, cv::COLOR_BGR2RGB);		// BGR->RGB
	cv::Mat image_transfomed;
	cv::resize(image, image_transfomed, cv::Size(256, 256));

	// Mat转Tensor
	torch::Tensor tensor_image = torch::from_blob(image_transfomed.data, { image_transfomed.rows, image_transfomed.cols, 3 }, torch::kByte);	//{ 256,256,3 }
	tensor_image = tensor_image.toType(torch::kFloat);					// 为了下一步归一化除255,将无符号整型转为float型
	tensor_image = tensor_image.div(255.0);								// 归一化
	tensor_image = tensor_image.permute({ 2,0,1 }).unsqueeze(0);		// opencv读取的图像矩阵存储形式:H x W x C, 但是Pytorch中Tensor的存储为:N x C x H x W, 因此需要进行变换
	tensor_image = tensor_image.to(torch::kCUDA);

	// perdict
	at::Tensor output = module.forward({ tensor_image }).toTensor();	// output.size = { 1,1,256,256 }
	at::Tensor predict(output);
	predict = predict.squeeze(0);										// 降维
	predict = predict.mul(255).clamp(0, 255).to(torch::kU8);			// 反归一化并转换为无符号整型			
	predict = predict.to(torch::kCPU);									// kCUDA报错

	// Tensor转Mat
	cv::Mat result(image_transfomed.rows, image_transfomed.cols, CV_8UC1, predict.data_ptr());	// (1) CV_8UC1√  CV_8UC3× (2) 必须是256×256,若是直接resize为原图大小则出错
	cv::resize(result, result, cv::Size(img_w, img_h));											// 还原为原尺寸
	cv::imshow("result", result);
	cv::imwrite("../prediction.png", result);
	cv::waitKey(0);

	return 0;
}

输入:

limesurvey部署 libtorch部署_opencv_10


输出:

limesurvey部署 libtorch部署_深度学习_11


done!!!

可能会遇到的问题:at::Scalar不明确
原因:使用Libtorch构建C++推理程序中,因为使用 opencv的Scalar类型,导致和Libtorch命名空间的Scalar冲突。
修改 libtorch/include/ATen/detail/CUDAHooksInterface.h的第26行:

namespace at{
using c10:Allocator;  // 添加命名空间

修改libtorch/include/ATen/core/TensorBody.h的第35行

namespace at{
using c10::Scalar;  //添加命名空间

参考:libtorch 中Scalar 歧义性

萌新一个,因为项目需要用C++,只有在没有Pytorch基础的情况下强上Libtorch,顺带也学习了一下Pytorch。Libtorch这部分的资料比较少,踩了很多坑,在此记录一下。

参考链接

[1] 在C++部署Pytorch(Libtorch)模型的方法总结(Win10+VS2017) [2] TorchScript使用的注意事项和常见错误 [3] C++部署Pytorch(Libtorch)出现问题、错误汇总 [4] libtorch 常用api函数示例(史上最全、最详细) [5] libtorch知识总结 [6] libtorch 中Scalar 歧义性