pytorch算法模型转换为onnx--＞ncnn（落地项目全生命周期3)

这一章节我们要对我们已经训练好的模型进行转换，分为三个步骤：1）pytorch模型转为onnx模型；2）onnx模型精简；3）onnx模型转为ncnn模型。

一：pytorch模型转onnx模型

我们以pytoch自带的算法为案例：

import torch
import torchvision

#define resnet18 model
model = torchvision.models.resnet18(pretrained=True)
#define input shape
x = torch.rand(1, 3, 224, 224) #N,C, W,H替换为你自己图片的大小
#define input and output nodes, can be customized
input_names = ["x"]
output_names = ["y"]
#convert pytorch to onnx
torch_out = torch.onnx.export(model, x, "resnet18.onnx", input_names=input_names, output_names=output_names,
 #dynamic_axes字段用于批处理.若不想支持批处理或固定批处理大小,移除dynamic_axes字段即可.
                        #'input0':[0, 2, 3]中数字0，2，3是指张量的维度，表示0，2，3维度可以动态尺寸输入。
                        # dynamic_axes={"input":{0:"batch_size"},  # 批处理变量
                        #                 "output":{0:"batch_size"}})

经过上面的步骤，我们得到了resnet18.onnx模型。然后我们测试下原始pytorch模型和onnx模型的推理结果是否一致，其代码如下所示：

import torch
import torchvision
import onnxruntime as rt
import numpy as np
import cv2

#test image
img_path = "test_onnx.jpg"
img = cv2.imread(img_path)
img = cv2.resize(img, (224, 224))
img = np.transpose(img, (2, 0, 1)).astype(np.float32)
img = torch.from_numpy(img)
img = img.unsqueeze(0)

#pytorch test
model = torchvision.models.resnet18(pretrained=True)
model.eval()
output = model.forward(img)
val, cls = torch.max(output.data, 1)
print("[pytorch]--->predicted class:", cls.item())
print("[pytorch]--->predicted value:", val.item())

#onnx test
sess = rt.InferenceSession("resnet18.onnx")
x = "x"
y = ["y"]
output = sess.run(y, {x : img.numpy()})
cls = np.argmax(output[0][0], axis=0)
val = output[0][0][cls]
print("[onnx]--->predicted class:", cls)
print("[onnx]--->predicted value:", val)

将test_onnx图片换做你自己的图片即可。

[pytorch]--->predicted class: 588
[pytorch]--->predicted value: 154.0394287109375
[onnx]--->predicted class: 588
[onnx]--->predicted value: 154.03944

由结果可知，算法经过转换后推理高度相似。我们的模型可以用这个网址来查看。
在python环境中：python server方式
step1、执行pip install netron安装包

step2、代码实现：

import netron
 
netron.start("model.onnx")

默认会在浏览器http://localhost:8080地址打开。若要修改端口，添加参数address=端口号即可。

二、精简onnx模型

有时候我们将得到的onnx模型转ncnn的时候，会报错，或者有些层不会被编译，这样我们可以将我们得到的onnx模型在精简些。
首先安装onnxsim包：

pip install onnx-simplifier

然后加载我们刚才得到的onnx文件，使用simplify函数对其进行精简处理：

from onnxsim import simplify
onnx_model = onnx.load(output_path)  # load onnx model
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, output_path)
print('finished exporting onnx')

然后即可得到的resnet18-sim.onnx（这个名字放在output_path中）模型。

三、onnx-sim模型转化为ncnn模型

我们在上一篇文章讲了ncnn工具的编译。我们基于我们编译好的ncnn，找到tools文件夹内的onnx2ncnn,该工具可以将onnx模型转为ncnn模型。命令：

onnx2ncnn resnet18-sim.onnx resnet18.param resnet18.bin

生成resnet18的param文件和bin文件，其中，param文件保存了模型结构。
打开param文件即可发现，模型的输入为x，输出为y，和第一步转onnx模型时定义的input_names = [“x”] 和 output_names = [“y”] 保持一致。若模型存在多个输入或多个输出，可自定义其输入和输出节点列表。
接下来我们测试下ncnn模型的推理结果的正确性，基于我们上一篇文章配置好的环境下，其代码如下：

#include <opencv2/highgui/highgui.hpp>
#include <vector>
#include "net.h"

using namespace std;
// m中存放着模型处理的结果
//这个函数作用对模型处理后得到的结果进一步分析
vector<float> get_output(const ncnn::Mat& m)  
{
	vector<float> res;
	for (int q = 0; q<m.c; q++)
	{
		const float* ptr = m.channel(q);
		for (int y = 0; y<m.h; y++)
		{
			for (int x = 0; x<m.w; x++)
			{
				res.push_back(ptr[x]);
			}
			ptr += m.w;
		}
	}
	return res;
}

int main()
{
	cv::Mat img = cv::imread("test_onnx.jpg");
	int w = img.cols;
	int h = img.rows;
	ncnn::Mat in = ncnn::Mat::from_pixels_resize(img.data, ncnn::Mat::PIXEL_BGR, w, h, 224, 224);
	
	ncnn::Net net;
	net.load_param("resnet18.param");
	net.load_model("resnet18.bin");
	ncnn::Extractor ex = net.create_extractor(); //初始化
	ex.set_light_mode(true);
	ex.set_num_threads(4);

	ex.input("x", in);
	ncnn::Mat feat;
	ex.extract("y", feat);  //使用模型处理图像，并得到结果
	vector<float> res = get_output(feat);
	vector<float>::iterator max_id = max_element(res.begin(), res.end());
	printf("predicted class: %d, predicted value: %f", max_id - res.begin(), res[max_id - res.begin()]);
	net.clear();
	return 0;
}

运行结果如下：

predicted class: 588, predicted value: 154.039322

预测类别和pytorch/onnx保持一致，由于计算库的不同，预测概率略微偏差。

四、ncnn模型的加载

我们在第二步得到了.parm文件和.bin文件。其中parm文件是模型文件，存放着算法模型，bin文件存放参数。因为parm文件是明文的，为了算法的保密性，不适合直接发布，因此我们将param文件转为二进制文件。在编译好的ncnn中的tools文件夹中，有个ncnn2mem工具，我们使用这个工具将生成的.parm转为.bin、.id.h、.mem.h三个文件，命令：

ncnn2mem resnet18.param resnet18.bin resnet18.id.h resnet18.mem.h

名字随便取，但是后缀不要变。其中生成的.param.bin文件是二进制模型的结构文件；.id.h文件是模型结构的头文件；.mem.h文件是模型参数的头文件。
我们为了将我们的算法和调用代码一起打包为dll文件，我们采用从内存加载parm和bin文件的方式。
截取部分resnet18.id.h文件分析：

namespace resnet18_param_id {
const int LAYER_x = 0;
const int BLOB_x = 0;   // 这是模型的输入参数
        .
        .
        .
const int LAYER_y = 77;
const int BLOB_y = 85;       //这是模型的输出参数，
} // namespace resnet18_param_id

从内存中加载模型的代码如下：

#include "resnet18.mem.h"
ncnn::Net net;
net.load_param(resnet18_param_bin);
net.load_model(resnet18_bin);

定义输入输出的代码如下：

#include "resnet18.id.h"
ncnn::Mat in;
ncnn::Mat out;
ncnn::Extractor ex = net.create_extractor();
ex.set_light_mode(true);
ex.set_num_threads(4);
ex.input(resnet18_param_id::BLOB_x, in);  //这里的BLOB_x就是.id.h文件中的输入
ex.extract(resnet18_param_id::BLOB_y, out); //这里的BLOB_y就是.id.h文件中的输出。这两个地方一定要对应！
ex.extract(resnet18_param_id::LAYER_y , out) //获得LAYER_y的输出结果

重点来了:ex.extract()可以针对某一层进行抽取得到结果。即你想看哪一层的输出，可以将BLOB_y更换为那一层的名字。如果你想要获取多个网络层的输出结果，这里可以多次写extract()方法。而且这样写并不会对模型的运行次数产生影响，即运行一次就能得到你想要的多个结果
整体的调用ncnn模型代码如下：

#include <opencv2/highgui/highgui.hpp>
#include <vector>
#include "net.h"
#include "resnet18.id.h"
#include "resnet18.mem.h"

using namespace std;

int main()
{
	cv::Mat img = cv::imread("test_onnx.jpg"); #换为你自己的
	int w = img.cols;
	int h = img.rows;
	ncnn::Mat in = ncnn::Mat::from_pixels_resize(img.data, ncnn::Mat::PIXEL_BGR, w, h, 224, 224);  # w,h换为你自己的
	
	ncnn::Net net;
	net.load_param(resnet18_param_bin);  #换为你自己的
	net.load_model(resnet18_bin);  #换为你自己的
	ncnn::Extractor ex = net.create_extractor();
	ex.set_light_mode(true);
	ex.set_num_threads(4);

	ncnn::Mat out;
	ex.input(resnet18_param_id::BLOB_x, in); #输入参数名称换为你自己的
	ex.extract(resnet18_param_id::BLOB_y, out);  #输出参数名称换为你自己

	ncnn::Mat out_flattened = out.reshape(out.w * out.h * out.c);
	vector<float> score;
	score.resize(out_flattened.w);
	for (int i = 0; i < out_flattened.w; ++i) {
		score[i] = out_flattened[i];
	}
	vector<float>::iterator max_id = max_element(score.begin(), score.end());
	printf("predicted class: %d, predicted value: %f", max_id - score.begin(), score[max_id - score.begin()]);

	net.clear();
	return 0;
}

注意：#不是c++代码的注释符号，这里的目的是让你知道哪些需要你自己替换为你自己的东西。
1）直接加载ncnn模型可以快速测试模型效果，但是param文件是明文的，打开文件可以直接看到模型结构，工程发布；
2）将param文件转为二进制，可以起到一定的加密作用；
3）将模型结构和参数直接读进内存，和代码整合在一起，利于打包发布，只需提供一个打包好的库即可（.so/.dll等），不用将模型文件单独拷贝到部署机器上，大大方便了算法的部署以及加密。