pytorch模型部署的几个方法

模型转换与调用

• 模型转换

• 转换torchscript格式
• 转换onnx格式

• 模型调用

• 利用opencv 载入并调用onnx格式模型
• 利用libtorch载入并调用pt格式模型

• 模型载入
• Mat to Tensor 类型转换
• 模型调用
• Tensor to Mat 类型转换
• Mat、Tensor数据格式总结
• c++输出txt

• Tensor 数据遍历
• Mat 数据遍历

模型转换

pytorch模型部署到c++需要转化为libtorch可读取的torchscript格式，或者更通用的onnx格式

转换torchscript格式

pytorch模型转换中用到trace方法运行一次网络并将模型中的参数与网络结构融合
注意当网络输出是list时trace中加入strict=False避免warning

net = PoseEstimationWithMobileNet() # 调用定义好的网络
    checkpoint = "./checkpoint_iter_370000.pth"
    net.load(checkpoint)
    net.eval（)

    forward_input = torch.rand(1, 3, 256, 394) # 模型的输入Tensor大小

    # trace_model = torch.jit.trace(net, forward_input, strict=False)  
    trace_model = torch.jit.trace(net, forward_input)  
    trace_model.save('./pose.pt')

转换onnx格式

net = PoseEstimationWithMobileNet()
    checkpoint = "./checkpoint_iter_370000.pth"
    net.load(checkpoint)
    net.eval（)

    forward_input = torch.rand(1, 3, 256, 394) # 模型的输入Tensor大小
    trace_model = torch.jit.trace(net, forward_input)
    # save onnx
    input_names = ["input"]
    output_names = ["output"]
    torch.onnx.export(net, forward_input, "pose.onnx", verbose=False, input_names=input_names,
                      output_names=output_names)

模型调用

利用opencv 载入并调用onnx格式模型

本文中利用opencv载入onnx
注意opencv的ReadNetFromTorch函数用来读取torch模型而不是pytorch,所以不能直接读取.pt文件

cv::dnn::Net inputNet = cv::dnn::readNet("pose.onnx");

调用模型

cv::Mat input = cv::imread(inputFile, cv::IMREAD_COLOR);
inputNet.setInput(input);

利用libtorch载入并调用pt格式模型

由于本项目中输入图片的预处理和输出图片的前向传播过程都使用opencv函数编写，格式也都是opencv的Mat格式，所以在Libtorch载入网络过程前后都需要进行libtorch中Tensor格式和opencv中Mat的互相转换.
这里提一下Tensor和Mat在获取高维矩阵结构时的不同

cv::Mat inputBlob = ...
int size0 = inputBlob.size[0]; //1 
int size1 = inputBlob.size[1]; //3
int size2 = inputBlob.size[2]; //256
int size3 = inputBlob.size[3]; //394

torch::Tensor img_tensor = ...
std::cout << "size:"<< img_tensor.sizes() << std::endl;

模型载入

torch::jit::script::Module module;
	module = torch::jit::load("/pose.pt" );

Mat to Tensor 类型转换

网上找到的代码大多为三维三通道图像通过torch::from_blob函数生成四维单通道Tensor，输入的维度信息大概如下

//opencv format H*W*C
 auto input_tensor = torch::from_blob(frame.data, {1, frame_h, frame_w, kCHANNELS});

然而本项目中需要处理图像经过预处理后生成的四维单通道Mat,经过多次数据输出对比，如下的代码可以成功转换数据格式并且不影响数据存放顺序

// torch::Tensor img_tensor = torch::from_blob(inputBlob.data, {inputBlob.size[0], inputBlob.size[1], inputBlob.size[2], inputBlob.size[3]},torch::kFloat32);
torch::Tensor img_tensor = torch::from_blob(inputBlob.data, {1,3,256,394},torch::kFloat32); //输入四维 和上一行代码等效 即按输入Mat维度顺序传入函数 不需要下一行的维度转换
	// img_tensor = img_tensor.permute({0,3,1,2});
  // img_tensor = img_tensor.div(255);  预处理时经过归一化处理 这里删掉

References :
https://zhuanlan.zhihu.com/p/92090874

模型调用

auto Output = module.forward({img_tensor}).toTensorList();

Tensor to Mat 类型转换

// netOutput = netOutput.squeeze().detach().permute({1, 2, 0}); 需要4维Mat所以不进行squeeze操作
// netOutput = netOutput.mul(255).clamp(0, 255).to(torch::kU8);网络需要Float结构不转ku8
int size[4];//创建Tensor结构大小的size数组 
for(int i=0;i<sizeof(netOutput.sizes())/sizeof(int);i++)
	size[i] = netOutput.sizes()[i] ;
cv::Mat netOutputBlob(4,size,CV_32FC1,Output.data_ptr());	
//注意需要的数据类型，这里选择CV_32FC1格式

References :
高维Mat的创建

Mat、Tensor数据格式总结

cv::Mat
 数值                     类型                          
CV_8U                 8位无符号整数                          
CV_8S                 8 位符号整数                         
CV_16U                16 位无符号整数                       
CV_16S                16 位符号整数                        
CV_32S                32 位符号整数                  
CV_32F                32 位浮点数                       
CV_64F                64 位浮点数

Mat_<uchar>---------CV_8U
Mat<char>-----------CV_8S
Mat_<short>---------CV_16S
Mat_<ushort>--------CV_16U
Mat_<int>-----------CV_32S
Mat_<float>----------CV_32F
Mat_<double>--------CV_64F

References :

c++输出txt

由于数据类型经过多次转换，数据的存储顺序可能发生变化，导致处理后的现象不正确，并且vscode调试功能并不方便，所以采用了较为直接的方法将转换前后数据输出txt进行对比

#include<iostream>
#include <fstream>
std::ofstream out("out.txt");//定义输出
out<<netOutput<<std::endl;
out.close();//关闭文件

References :

Tensor 数据遍历

Tensor和Mat数据遍历格式不同，输出代码也有所不同
Tensor可以直接输出

test<<netOutput<<std::endl;//tensor输出txt

Mat 数据遍历

for(int i=0;i<netOutputBlob.size[1];i++) //mat遍历输出保存txt
	{for(int j=0;j<netOutputBlob.size[2];j++)
		{for(int k=0;k<netOutputBlob.size[3];k++)
			in<<netOutputBlob.ptr<float>(0,i,j)[k]<<std::endl;// 输出类型要根据Mat数据类型改变 比如float int 参考前文的Mat数据类型总结
 		}
}