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