欢迎您访问 最编程 本站为您分享编程语言代码,编程技术文章!
您现在的位置是: 首页

pytorch 到 ncnn 及其测试 - I. pytorch 模型到 onnx 模型

最编程 2024-03-28 12:08:24
...

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