pytorch转ncnn及其测试

       本文主要讲解pytorch转ncnn的过程，以及其推理结果的正确性测试。

       其中，模型转换分为如下两步：1. pytorch模型转onnx模型；2. onnx模型转ncnn模型。

       onnx(https://github.com/onnx/onnx)是一种开放神经网络交换的格式，可使模型在不同框架之间进行转换。

       一、pytorch模型转onnx模型

       采用pytorch自带的resnet18作为示例，代码如下：

import torch
import torchvision

#define resnet18 model
model = torchvision.models.resnet18(pretrained=True)
#define input shape
x = torch.rand(1, 3, 224, 224)
#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)

       至此，得到转换后的resnet18.onnx模型。

       测试pytorch模型和onnx模型的推理结果是否一致，代码如下：

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

#test image
img_path = "test.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.jpg换成自己的图片即可，测试结果如下：

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

        推理结果是对的，预测概率会有些许的偏差。

       二、onnx模型转ncnn模型

       ncnn(https://github.com/Tencent/ncnn)编译好后，tools内有一个onnx2ncnn的工具，直接使用此工具即可将onnx模型转换为ncnn模型，命令如下：

onnx2ncnn resnet18.onnx resnet18.param resnet18.bin

       生成resnet18的param文件和bin文件，其中，param文件保存了模型结构，bin文件保存了模型参数。

       打开param文件即可发现，模型的输入为x，输出为y，和第一步转onnx模型时定义的input_names = [“x”] 和 output_names = [“y”] 保持一致。若模型存在多个输入或多个输出，可自定义其输入和输出节点列表。

       测试ncnn模型的前向推理结果的正确性，导入ncnn的库和头文件后，调用代码如下：

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

using namespace std;

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.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保持一致，由于计算库的不同，预测概率略微偏差。