图像处理深度学习python代码改写成c++推理

目录

前言

一、模型导出与导入

 导出模型

加载模型

一、代码改写

 ONNX模型

总结


前言

深度学习的代码一般采用python编写,因为python有很多开源的库可以直接引用,但是如果对推理速度有更高的要求,可以采用c++编写推理代码。相比python,c++运行速度更加快,但是需要编译环境,而且第三方的库没有python完善。所以,究竟要使用哪一种语言,可以根据自己的应用场景和使用目的来选择。

如果我们需要用c++来推理,而源代码又只有python版本,那么就需要我们来改写。首先我们需要在python代码中将模型训练好,或者使用预训练模型。接下来我将从两个方面讲述移植的过程。

一、模型导出与导入

python训练出来的模型是没有办法直接在c++上使用的,所以我们需要对模型做一定的处理。对于pytorch框架的模型,c++有libtorch。libtorch是pytorch的C++版本,支持CPU端和GPU端的部署和训练。由于python和c++的语言特性,因此用pytorch做模型训练,libtorch做模型部署。首先我们需要在python代码中将模型用litorch导出,然后在c++代码中加载运行。

图像处理深度学习python代码改写成c++推理_第1张图片

 导出模型

导出模型的代码一定要在加载模型权重之后,不然导出的模型是没有权重参数的。导出代码如下:

## 导出模型
model_ = torch.jit.trace(model.eval(), torch.randn(1, 3, 256, 256))
torch.jit.save(model_,'model.pt')

torch.jit.trace函数的第一个参数是需要导出的模型(torch.nn.Model),第二个参数是模型的输入(torch.Tensor)。导出的原理是跟踪模型输入,记录所有执行到图形中的操作。输入张量的内容可以是随机的,但是尺寸必须符合模型规定的输入大小。

torch.jit.save函数的第一个参数是torch.jit.trace函数的输出,第二个参数是保存模型的路径及名字。

运行成功后,将会在指定路径下生成一个pt文件,这就是我们需要的模型。

加载模型

在c++代码中,我们将导出的模型加载推理。加载代码如下:

// 加载模型
model = torch::jit::load("model_.pt");
model.to(torch::kCUDA);  //加载到gpu运行

可以指定加载的模型在cpu还是gpu运行,如果模型加载到gpu上,那么之后输入模型的数据也需放在gpu上。

二、代码改写

对于图像的处理,python代码中可能会使用opencv、numpy、PIL、skimage等库,但是在c++中我们没有那么多开源的库可以使用,而且配置这些库也比较麻烦。其实对于图像的操作,我们基本上可以用opencv来实现,我们需要做的就是,看懂python代码对图像或者张量的处理,然后用opencv将其实现。所以我们的c++代码运行环境必须配备opencv。

opencv表示图像的数据类型是cv::Mat,而模型输入需要的数据类型是torch::Tensor,所以我们最后需要转化数据类型。整体的数据流如下图所示:

图像处理深度学习python代码改写成c++推理_第2张图片

根据不同代码的图片预处理和后处理,我们需要写出对应的opencv代码。下面我举例几个图像或者张量的操作。

cv:Mat image,input,output;

// 读图片
image = cv::imread("1.jpg");

// 8U 转 32F
input.convertTo(output, CV_32FC3, 1.0 / 255.);

// 32F 转 8U
input.convertTo(output, CV_8UC3,255);

// 调整图片大小
cv::resize(input, output, cv::Size(H, W));

// 图片颜色描述空间转换
cv::cvtColor(input, output, cv::COLOR_BGR2Lab);  //opencv一开始读进来的图片是BGR格式的
cv::cvtColor(input, output, cv::COLOR_BGR2RGB);
cv::cvtColor(input, output, cv::COLOR_BGR2GRAY);

// 截取某一个通道
cv::extractChannel(input, output, 0);
cv::extractChannel(input, output, 1);
cv::extractChannel(input, output, 2);

// 拆分通道
std::vector Channels(3);
cv::split(input, Channels);

// 合并通道
std::vector Channels;
Channels.push_back(通道0);
Channels.push_back(通道1);
Channels.push_back(通道2);
cv::merge(Channels,output);

// 保存图片
cv::imwrite("1.jpg", image);

另外,还涉及到cv::Mat和torch::Tensor互相转换的问题,这里我提供两个转换函数仅供参考。

cv::Mat tensor2mat(torch::Tensor tensor, bool show = false)
{
    torch::Tensor ten=tensor.to(torch::kCPU);
    int H, W;
    H = ten.size(2);
    W = ten.size(3);
    int channel = ten.size(1);
    int index = 0;
    float *rst_data = ten.data_ptr<float>();
    cv::Mat ret = cv::Mat::zeros(H, W, CV_32FC3); //三通道
    
    // 不同tensor的排列方式不同,请根据实际排列方式修改下面代码
    // ============================================================
    for (int h = 0; h < H; h++)
    {
        for (int w = 0; w < W; w++)
        {
            for (int c = 0; c < channel; c++)
            {
                float val = rst_data[index];
                ret.at(h, w)[c] = val;
                index++;
            }
        }
    }
    // ==============================================================

    if (show) // 打印结果
    {
        cout << endl
             << "===================== Tensor ============================="
             << endl;
        cout << endl << ten;
        cout << endl
             << "===================== Mat ============================="
             << endl;
        cout << ret << endl;
    }
    return ret;
}
torch::Tensor mat2tensor(cv::Mat mat, bool show = false)
{
    torch::Tensor ret;
    ret = torch::from_blob(mat.data,
                           {1, mat.rows, mat.cols, mat.channels()},
                           torch::TensorOptions().dtype(torch::kFloat32));
    ret = ret.permute({0, 3, 1, 2}).to(torch::kCUDA);

    if (show) // 打印结果
    {
        cout << endl
             << "===================== Mat ============================="
             << endl;
        cout << endl << mat;
        cout << endl
             << "======================Tensor ============================"
             << endl;
        cout << ret << endl;
    }
    return ret;
}

c++模型推理代码:

// 推理
std::vector input_tensor = {输入的tensor}; 
torch::Tensor output = model.forward(input_tensor).toTensor();

 ONNX模型

以上主要针对于pytorch框架的移植,如果使用的是其他框架,则可以使用ONNX模型。无论你使用什么样的训练框架来训练模型(比如TensorFlow/Pytorch/OneFlow/Paddle),你都可以在训练后将这些框架的模型统一转为ONNX存储。 ONNX文件不仅存储了神经网络模型的权重,还存储了模型的结构信息、网络中各层的输入输出等一些信息。 然后将转换后的ONNX模型,转换成我们需要使用不同框架部署的类型,通俗来说ONNX 相当于一个翻译。

如果使用ONNX模型,则需修改模型导出代码。ONNX模型可以使用opencv加载运行,需要注意的是,用opencv加载模型后,模型的输入输出就不再是Tensor,而是Mat了。这里的Mat需要使用cv::dnn::blobFromImage()函数进行转化。

opencv加载ONNX模型相关函数具体详见:函数详解

其他相关教程详见:ONNX相关教程

举例将pytorch模型导出为ONNX格式的代码:

# 导出onnx模型
torch.onnx.export(model, 
 		  		torch.randn(1, 3, 256, 256), # 模型输入
            	"siggraph17.onnx",   # 模型名字
 			    export_params=True,  # 导出模型参数
 				input_names=['input'], 
				output_names=['output'], 
				dynamic_axes={'input':{0:'N', 2:'H', 3:'W'},'output':{0:'N', 2:'H',3:'W'}},                  # 设置动态的维度 
				opset_version=11)

 c++模型推理代码:

model.setInput(input);
cv::Mat output = model.forward();

总结

以上就是我在python->c++中的经验和总结,如有问题望指正,欢迎找我探讨。

你可能感兴趣的:(python,深度学习,c++,pytorch)