pytorch resnet18转换为ncnn

模型简化:

python3 -m onnxsim input_onnx_model output_onnx_model

转换完成后,需要注意:
1.图像归一化方式:
如果训练过程中,图像归一化为:

preprocess_transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
        ])

那么,ncnn测试时,同样需要讲图像归一化:

ncnn::Mat in = ncnn::Mat::from_pixels_resize(image.data, ncnn::Mat::PIXEL_BGR2RGB, w, h, 224, 224);
const float mean_vals[3] = {0.485f*255.f, 0.456f*255.f, 0.406f*255.f};
const float norm_vals[3] = {1/0.229f/255.f, 1/0.224f/255.f, 1/0.225f/255.f};
in.substract_mean_normalize(mean_vals, norm_vals);

其中ncnn::Mat::PIXEL_BGR2RGB,根据训练的格式变动。

完整代码如下

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

void resnet_init(ncnn::Net* net, const char* paramfile, const char* binfile)
{
	if(net == NULL) net = new ncnn::Net;
	net->load_param(paramfile);
	net->load_model(binfile);
}

void resnet_release(ncnn::Net* net)
{
	if(net == NULL) return;
	net->clear();
	delete net;
}

static int print_topk(const std::vector<float>& cls_scores, int topk)
{
    // partial sort topk with index
    int size = cls_scores.size();
    std::vector<std::pair<float, int> > vec;
    vec.resize(size);
    for (int i = 0; i < size; i++)
    {
        vec[i] = std::make_pair(cls_scores[i], i);
    }

    std::partial_sort(vec.begin(), vec.begin() + topk, vec.end(),
                      std::greater<std::pair<float, int> >());

    // print topk and score
    for (int i = 0; i < topk; i++)
    {
        float score = vec[i].first;
        int index = vec[i].second;
        fprintf(stderr, "%d = %f\n", index, score);
    }

    return 0;
}

void softmax(std::vector<float>& scores)
{
	float sum = 0.0f;
	for(int i = 0; i < scores.size(); i++)
	{
		float s = scores[i];
		scores[i] = std::exp(s);
		//scores[i] = std::exp(scores[i]);
		sum += scores[i];
	}
	printf("sum = %.3f\n", sum);
	for(int i = 0; i < scores.size(); i++)
	{
		scores[i] /= sum;
	}
}

void findmax(std::vector<float> scores)
{
	float max = -1;
	int index = -1;
	for(int i = 0; i < scores.size(); i++)
	{
		if(scores[i] > max){
			max = scores[i];
			index = i;
		}
	}
	printf("max: %.3f, id = %d\n", max, index);
}

void forward(ncnn::Net* net, const cv::Mat& image)
{
	int w = image.cols;
	int h = image.rows;
	ncnn::Mat in = ncnn::Mat::from_pixels_resize(image.data, ncnn::Mat::PIXEL_BGR2RGB, w, h, 224, 224);

    const float mean_vals[3] = {0.485f*255.f, 0.456f*255.f, 0.406f*255.f};
	const float norm_vals[3] = {1/0.229f/255.f, 1/0.224f/255.f, 1/0.225f/255.f};
	in.substract_mean_normalize(mean_vals, norm_vals);

	
	ncnn::Mat out;
	ncnn::Extractor ex = net->create_extractor();
	ex.set_light_mode(true);
	ex.set_num_threads(4);
	ex.input("label", in);
	ex.extract("result", out);


    ncnn::Layer* softmax = ncnn::create_layer("Softmax");

    ncnn::ParamDict pd;
    softmax->load_param(pd);

    softmax->forward_inplace(out, net->opt);

    delete softmax;

    out = out.reshape(out.w * out.h * out.c);
    std::vector<float> cls_scores;
    cls_scores.resize(out.w);
    for (int j = 0; j < out.w; j++)
    {
        cls_scores[j] = out[j];
    }
    print_topk(cls_scores, 2);

}

int main(int argc, char** argv)
{
	ncnn::Net* net = new ncnn::Net;

	resnet_init(net, "models/resnet18_cat_dog.param", "models/resnet18_cat_dog.bin");
	cv::Mat src = cv::imread("cat_dog_test/dog/33.jpg");
	forward(net, src);
	delete net;
	return 0;
}
posted @ 2020-12-18 18:31  快乐码小农  阅读(839)  评论(0编辑  收藏  举报