C++ onnxruntime 推理例子

C++ onnxruntime 推理例子

flyfish

假如一个模型的输入和输出如图所示

输入节点名字是 input，形状是[batch_size,3,128,64]
输出节点名字是output,形状是 [batch_size,512]

源码如下

#include <assert.h>
#include <vector>
#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>
#include "cpu_provider_factory.h"
#include <iostream>


int main(int argc, char* argv[]) {
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "Default");

    const char* model_path = "./model.onnx";



    Ort::Session session_{env, model_path, Ort::SessionOptions{nullptr}}; //CPU


    static constexpr const int height_ = 128; //model input height
    static constexpr const int width_ = 64; //model input width

    Ort::Value input_tensor_{nullptr};
    std::array<int64_t, 4> input_shape_{1, 3, height_, width_}; //mode input shape NCHW = 1x3xHxW

    Ort::Value output_tensor_{nullptr};
    std::array<int64_t, 2> output_shape_{1, 512}; //model output shape,

    std::array<float, width_ * height_ * 3> input_image_{};
    std::array<float, 512> results_{};

    //rows = height 128
    //cols = weight 64

    std::string imgPath = "./1.jpg";
    cv::Mat img = cv::imread(imgPath);
    cv::resize(img, img, cv::Size(64, 128));//size
    cv::cvtColor(img, img, cv::COLOR_BGR2RGB);//bgr -> rgb
    std::cout<<"dim:"<< img.rows<<":"<<img.cols<<":"<<img.channels()<<std::endl;
    //根据自己的处理更改
    int index=0;
    for (unsigned c = 0; c < 3; c++)
    {
        for (unsigned i = 0; i < img.rows; i++)
        {
            for (unsigned j = 0; j < img.cols; j++)
            {
                input_image_[index++] = (static_cast<float>(img.at<cv::Vec3b>(i, j)[c]));
            }
        }
    }


    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);
    input_tensor_ = Ort::Value::CreateTensor<float>(memory_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());
    output_tensor_ = Ort::Value::CreateTensor<float>(memory_info, results_.data(), results_.size(), output_shape_.data(), output_shape_.size());
    const char* input_names[] = {"input"};
    const char* output_names[] = {"output"};


    session_.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor_, 1, output_names, &output_tensor_, 1);


    float* out = output_tensor_.GetTensorMutableData<float>();
    for(int i = 0; i < 512; i++)
    {
        std::cout<< out[i]<<":";
    }

    return 0;
}

CMakeLists.txt内容如下

cmake_minimum_required(VERSION 3.5)

project(onnx_example LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_STANDARD_REQUIRED ON)



set(ONNXRUNTIME_DIR "/libpath/onnxruntime-linux-x64-1.12.1")
include_directories("${ONNXRUNTIME_DIR}/include"
    )
find_package(OpenCV REQUIRED )


include_directories(
    ${OpenCV_INCLUDE_DIRS}/include
    )


add_executable(onnx_example
main.cpp)
target_link_libraries(onnx_example PRIVATE "${ONNXRUNTIME_DIR}/lib/libonnxruntime.so" ${OpenCV_LIBS} )