AODNet在Libtorch端的部署全过程(附详细源码)

目录

一、部署环境

二、部署过程

1、VS2017配置Opencv的环境

2、CMake的安装

3、PyTorch —> Libtorch模型的转换

4、Libtorch端的实现

5、遇到的问题

---

一、部署环境

Windows10、Opencv3.4.10、Libtorch1.5.0-CPU-Release、PyTorch1.5.0-CUDA10.1、VS2017、CMake3.19

二、部署过程

1、VS2017配置Opencv的环境

参考博客。感谢这位博主的分享，亲测有效！

2、CMake的安装

这个安装比较简单，网上很多教程，注意中间勾选上添加到环境变量，其他的傻瓜式一键到底安装就行。

3、PyTorch —> Libtorch模型的转换

使用Libtorch在Windows端对AODNet进行部署的流程主要分为两部分：一个是将在python端训练好的模型文件(*.pth文件)转换为Libtorch使用的模型文件(*.pt文件)。另外一部分是在Libtorch框架下使用C++对第一步生成的模型进行调用和前向传播，并重写前行传播的pipline，实现模型功能的C++实现。

这里使用trace方式将PyTorch模型转换为Libtorch模型(需要有提前训练好的PyTorch模型)。源码如下：

# pytorch2libtorch.py
import torchvision.models as models
from torchvision import transforms as transform
import torch
from net import dehaze_net
# 载入模型
model = dehaze_net()
# 载入权重
model.load_state_dict(torch.load(r"F:\\3-code\AODNet\\PyTorch-Image-Dehazing\\snapshots\\dehazer.pth", map_location=torch.device('cpu')))
# 使用eval进入predict阶段
model.eval()
# 随机初始化一个输入，用以通过模型进行记录
sample = torch.rand(1, 3, 640, 480)
# 跟踪并记录模型
traced_script_module = torch.jit.trace(model, sample)
# 查看输出的信息，该步骤仅验证使用，可不要
output = traced_script_module(sample)
print(output.size())
# 保存*.pt文件，即C++/libtorch需要调用的模型文件
traced_script_module.save("F:\\3-code\\AODNet\\AODNet-Libtorch\\dehaze_net.pt")

跟trace方式的字面意思一样，就是随机初始化一个输入，让其在装载了权重的网络中跑一边，在这过程中记录下模型的样子和参数等属性，然后将该记录另存为*.pt文件，即Libtorch的模型文件（可以理解为山寨版但效果一样的模型吧）。

4、Libtorch端的实现

这部分是整个过程中的重点，我们使用CMake进行编译，如果没用过CMake，千万不要被吓到，CMake其实很简单的，用一次就知道，相比在VS中手动配置一堆乱七八糟的东西，还不如直接上CMake香（用了CMake就不用在在VS里手动配置一堆环境了），真的！另外，建议Libtorch与PyTorch的版本是一样的，虽然确实有些情况下还是兼容的，但相同版本的真的能避免出现一些谜之问题（血泪教训！）。

首先建立一个空文件夹，在该文件夹里新建CMakelists.txt文本文件、example.cpp文件、build空文件夹，另外为了好找我把我生成的*.pt文件也塞进来了，一共是四个东西，见图1。先不要管都是干啥的，后面会逐个详细解说。

图1

首先是CMakelists.txt文本文件的编写说明：

先看内容如下，注释都写在里面了，有关语句想深入了解的话可以百度下。CMakelists文件主要定义了文件编译需要的依赖库以及一些编译规则等。

（一些适合CMake入门的链接：CMake 入门实战 | HaHack，软件构建: CMake 快速入门_哔哩哔哩_bilibili）

# 这里定义了 CMake 的最低版本号要求，此处最低版本要求是3.12
cmake_minimum_required(VERSION 3.12 FATAL_ERROR)
# 设置项目名称为AODNet，可自定义
project(AODNet)
# 添加依赖链接库的路径，这里需要添加opencv与libtorch的路径，用分号隔开。注意这里填自己解压libtorch时的路径
set(CMAKE_PREFIX_PATH "E:/opencv/opencv/build/x64/vc15/lib;E:/LibTorch/libtorch150cpu/libtorch") 

# 使用find_package引入外部依赖包
find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)

# 查询_FOUND变量，代表是否找到该依赖包
if(NOT Torch_FOUND)
    message(FATAL_ERROR "Pytorch Not Found!")
endif(NOT Torch_FOUND)
# 查询_LIBRARIES变量
message(STATUS "Pytorch status:")
message(STATUS "    libraries: ${TORCH_LIBRARIES}")

# 打印信息
message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    libraries: ${OpenCV_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")

# 指定生成可执行文件名为example
add_executable(example example.cpp)
# 指明可执行文件example需要连接的链接库
target_link_libraries(example ${TORCH_LIBRARIES} ${OpenCV_LIBS})
# 在给定的作用域内设置一个命名的属性
set_property(TARGET example PROPERTY CXX_STANDARD 11)

# 在命令行cd到build目录下，执行以下指令，后面两个“..”代表上级目录，因为相对build文件内来说，CMakelists文件在其上一级目录中
# cmake -DCMAKE_BUILD_TYPE=Release -G "Visual Studio 15 Win64" ..

然后是build文件夹的使用说明：

这个文件夹只是一个习惯的命名而已，如果愿意，名字可以随便取的。该文件就是我们存放编译产生文件的地方。具体用法见上方代码最后两行。这里说明下，我们用的Release模式进行编译，因此Libtorch库也需要使用Release版的。Visual Studio 15 代表VS2017，Win64代表64位。先别慌开始编译——我们还没说example.cpp文件呐！

下一个就是example.cpp文件的编写：

直接上全部的代码吧，其实很简单，就不一一解释了。注释写的很清楚了。

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

// tensor->Mat 格式转换
cv::Mat tensor2mat(torch::Tensor &tensor)
{
	tensor = tensor.squeeze().mul(255).add(0.5).clamp(0, 255).permute({ 1, 2, 0 }).to(torch::kU8);
	int height = tensor.size(0), width = tensor.size(1);
	// CV_8UC3含义：CV_(S|U|F)C 此处就是创建一个height * width的8位3通道的图像矩阵
	cv::Mat _2Mat(cv::Size(height, width), CV_8UC3, tensor.data_ptr());
	// 或者下面用这种方法
	//cv::Mat _2Mat(width, height, CV_8UC3);
	//std::memcpy((void *)_2Mat.data, tensor.data_ptr(), sizeof(torch::kU8)*tensor.numel());  // tensor.numel()：number of elements  返回元素个数
	cv::cvtColor(_2Mat, _2Mat, CV_RGB2BGR);
	return _2Mat;
}

// 生成去雾图像的保存地址
const char* gen_savePath(const char* image_path)
{
	std::string strPath = image_path;
	std::string prefix, suffix, savePath;
	int npos = strPath.find_last_of("/");
	if (npos != -1)
	{
		prefix = strPath.substr(0, npos);
		suffix = "/dehazed-" + strPath.substr(npos+1,  - 1);
		savePath = prefix + suffix;
	}
	std::cout << savePath << std::endl;                                                                                                                               
	return savePath.data();
}

// 去雾
void ImgDehaze(const char* image_path, const char* module_path)
{
	// 去雾图像的保存地址
	const char* savePath;
	// 读取图片,读取的是BGR格式的图片
	cv::Mat srcimage = cv::imread(image_path);
	cv::Mat image;
	// 图片格式转换为RGB
	cv::cvtColor(srcimage, image, CV_BGR2RGB);
	// 图像变换到480*640 
	// cv::resize(image, image, cv::Size(480, 640));
	// 将Mat类转化为tensor类型  opencv格式（h，w，c）
	at::Tensor img_tensor = torch::from_blob(image.data, { 1,  image.cols, image.rows, 3 }, torch::kByte);
	// 转换为torch::kFloat格式的数据
	img_tensor = img_tensor.toType(torch::kFloat32);

	// 转化为（b，c，h，w）格式
	img_tensor = img_tensor.permute({ 0, 3, 1, 2 });

	// 归一化
	img_tensor = img_tensor.div_(255.0);

	// 使用torch::jit::load()加载pt模型文件.
	torch::jit::script::Module module = torch::jit::load(module_path);
	//module.eval();

	// 加载输入
	std::vector input;
	input.push_back(img_tensor);
	// 将图像通过模型
	torch::Tensor output = module.forward(input).toTensor();
	// 将结果转化为Mat格式的图片
	cv::Mat resultImg = tensor2mat(output);
	// 生成去雾图像的保存地址
	savePath = gen_savePath(image_path);
	// 保存
	// std::vector compression_params;
	// compression_params.push_back(CV_IMWRITE_JPEG_QUALITY);  //选择jpeg
	// compression_params.push_back(100); //在这个填入你要的图片质量
	// cv::imwrite(savePath, resultImg, compression_params);
    cv::imwrite(savePath, resultImg);
	// 显示去雾后的图片
	cv::imshow("dehaze_img", resultImg);
	cv::waitKey(0);
}

//其中，argc = argument count ：表示传入main函数的数组元素个数，为int类型，
//而 argv = argument vector ：表示传入main函数的指针数组，为char**类型。
//第一个数组元素argv[0]是程序名称，并且包含程序所在的完整路径。argc至少为1，即argv数组至少包含程序名。
int main(int argc, char* argv[]) 
{
	// 检查参数个数
	//if (argc != 2)
	//{
		//std::cout << "arg num error,please input 2 args" << std::endl;
	//}
	const char* imgPath = "F:/3-code/AODNet/PyTorch-Image-Dehazing/test_images/test.png";
	const char* modulePath = "F:/3-code/AODNet/AODNet-Libtorch/dehaze_net.pt";

	ImgDehaze(imgPath, modulePath);

	while (1);
}

然后就是cmake编译的步骤

使用上面讲过的方法，进行编译，不出问题，挺快就好了，我这步没出啥大问题，注意有关路径要写正确，正常情况下会出现下面两行语句，代表编译通过了，如果cmake版本比较新的话编译过程中可能会有warning出现，可以不用理会。

-- Generating done
-- Build files have been written to: F:/3-code/AODNet/AODNet-Libtorch/build

好了，编译完成后我们可以打开build文件夹看一下，可以发现有同一个名为AODNet.sln的文件，是不是有点眼熟？没错，这个就是在CMakelists文件中设置的project(AODNet)的结果，还有一个example.vcxproj文件，同样的，我们在CMakelists中设置了可执行文件名为example，这个文件是这次的主角。

注：sln文件与vcxproj的区别与联系：sln是解决方案的配置，一个解决方案里面可以包含多个工程，主要是管理这个方案里的多个vcxproj，vcxproj是工程的配置文件，管理工程中细节比如包含的文件，引用库等，一般没有sln，也可以直接打开vcxproj，也可以重新生成sln，sln里有多个工程，当你移除某个工程时sln会有变化，sln并不是太重要

双击example.vcxproj文件就可以进入到VS2017中，我们在解决方案资源管理器中选中example，右键选择设为启动项目。可以发现example.cpp文件就在这里，如图2。这时我们选择Release及x64模式，点击本地Windows调试器（见图3）等待代码的运行。不出意外，系统会提示找不到c10.dll文件，没关系，我们找到Libtorch的解压文件夹，把lib目录下的dll文件全部放到build目录下的Release文件中，然后再次点击本地Windows调试器运行。不出问题就可以成功运行啦，若有问题欢迎留言交流。注意一点的是，libtorch版本最好与PyTorch版本对应，另外对于本例来说，libtorch要下载Release版的！

图2

图3

5、遇到的问题

我的运行结果如图4，可以看出网络去雾效果还是挺好的，问题就是在libtorch中运行之后回复出来的图片会莫名的多出几道白色的条纹，还挺有规律的。。。目前我还没找到bug所在T_T，有知道问题的小伙伴欢迎留言交流哈~

图4 去雾后

图5 去雾前