Libtorch部署显著性目标检测网络BASNet
- 一、环境配置
- 二、torchscript生成.pt文件
- 三、Libtorch部署
- 1.添加环境变量
- 2.VS配置环境
- 四、BASNet demo
- 参考链接
Libtorch 是Pytorch官方提供的C++ API,使用方法可以说是很大程度上还原了Pytorch。
一、环境配置
- GitHub地址:BASNet
- VS2017+CUDA 10.1+cuDNN 7.6.5+OpenCV 3.4.11
- Pytorch 1.5.0+torchvision 0.6.0
- Libtorch 1.5.0 (release):提取码:865k
https://pan.baidu.com/s/1Ty-UPZWEOnNRPwldexLZzw
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文件列表如下:
1.添加环境变量
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>
。
ii.项目—》属性—》链接器—》常规—》附加库目录
iii.项目—》属性—》链接器—》输入—》附加依赖项
不放心的话,可以把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检查 :否
v.项目—》属性—》C/C++目录—》语言—》符合模式 :否
vi.新建项目—》属性—》连接器—》命令行—》其他选项:
添加**/INCLUDE:?warp_size@cuda@at@@YAHXZ**
测试
#include "torch/torch.h"
#include "torch/script.h"
int main()
{
torch::Tensor output = torch::randn({ 3,2 });
std::cout << output;
return 0;
}
输出:
配置成功!
四、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;
}
输入:
输出:
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; //添加命名空间
萌新一个,因为项目需要用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 歧义性