MNN框架学习（四）：tensorflow图像分类模型部署

1.模型转换

先下载tensorflow的模型，下载地址为：

https://github.com/tensorflow/models/tree/master/research/slim

然后，使用编译好的MNN工具转换模型：

./MNNConvert -f TF --modelFile mobilenet_v1_1.0_224_frozen.pb --MNNModel mobilenet.mnn --bizCode MNN

2、模型部署

主要分为两个步骤：

第一步，初始化步骤，包括读取模型创建解释器，配置调度参数、配置后端参数和创建会话

int Classifier::Init(const char* root_path) {
    std::cout << "start Init." << std::endl;
    std::string model_file = std::string(root_path) + "/mobilenet.mnn";
    // 创建解释器
    classifier_interpreter_ = std::shared_ptr<MNN::Interpreter>(MNN::Interpreter::createFromFile(model_file.c_str()));
    
    if (!classifier_interpreter_ || LoadLabels(root_path) != 0) {
        std::cout << "load model failed." << std::endl;
        return 10000;
    }    
    
    // 配置调度
    MNN::ScheduleConfig schedule_config;
    schedule_config.type = MNN_FORWARD_CPU;
    schedule_config.numThread = 1;
    // 配置后端
    MNN::BackendConfig backend_config;
    backend_config.precision = MNN::BackendConfig::Precision_Normal;
    schedule_config.backendConfig = &backend_config;

    // 创建会话
    classifier_sess_ = classifier_interpreter_->createSession(schedule_config);
    input_tensor_ = classifier_interpreter_->getSessionInput(classifier_sess_, nullptr);

    classifier_interpreter_->resizeTensor(input_tensor_, {1, 3, inputSize_.height, inputSize_.width});
    classifier_interpreter_->resizeSession(classifier_sess_);

    std::cout << "End Init." << std::endl; 
    
    initialized_ = true;

    return 0;
}

第二步：数据读入、模型推理和后处理输出

int Classifier::Classify(const cv::Mat& img_src, std::vector<ImageInfo>* images) {
    std::cout << "start classify." << std::endl;
    images->clear();
    if (!initialized_) {
        std::cout << "model uninitialized." << std::endl;
        return 10000;
    }

    if (img_src.empty()) {
        std::cout << "input empty." << std::endl;
        return 10001;
    }

    cv::Mat img_resized;
    cv::resize(img_src.clone(), img_resized, inputSize_);
    std::shared_ptr<MNN::CV::ImageProcess> pretreat(
        MNN::CV::ImageProcess::create(MNN::CV::BGR, MNN::CV::RGB, meanVals, 3, normVals, 3)
    );
    pretreat->convert((uint8_t*)img_resized.data, inputSize_.width, inputSize_.height, img_resized.step[0], input_tensor_);

    // forward
    classifier_interpreter_->runSession(classifier_sess_);

    // get output
    // mobilenet: "classifierV1/Predictions/Reshape_1"
    MNN::Tensor* output_score = classifier_interpreter_->getSessionOutput(classifier_sess_, nullptr);

    // copy to host
    MNN::Tensor score_host(output_score, output_score->getDimensionType());
    output_score->copyToHostTensor(&score_host);

    auto score_ptr = score_host.host<float>();
    std::vector<std::pair<float, int>> scores;
    for (int i = 0; i < 1000; ++i) {
        float score = score_ptr[i];
        scores.push_back(std::make_pair(score, i));
    }

    std::partial_sort(scores.begin(), scores.begin() + topk_, scores.end(), std::greater< std::pair<float, int> >());
    for (int i = 0; i < topk_; ++i) {
        ImageInfo image_info;
        image_info.label_ = labels_[scores[i].second];
        image_info.score_ = scores[i].first;
        images->push_back(image_info);
    }

    std::cout << "end classify." << std::endl;

    return 0;
}

具体代码已经上传到github：https://github.com/MirrorYuChen/mnn_example/tree/master/src/classifier

大家觉得有用就给个star，不许白嫖哦~

参考资料：

[1] https://github.com/alibaba/MNN

[2] https://github.com/Linzaer/Ultra-Light-Fast-Generic-Face-Detector-1MB/tree/master/MNN