• 使用protobuf解析Onnx文件

    使用OpenCV加载Onnx推理的时候，无法获取到Onnx的网络输入大小，并且对推理速度要求不高不需要使用TensorRT的时候，如何才能得知Onnx的一些必要的信息，OpenCV没有提供接口，只能自己从Onnx文件中解析了。

    Onnx文件是使用protobuf序列化后的二进制数据，想要读取里面的信息需要使用protobuf将其反序列化为对象才行。

    第一步：

    编译 protobuf，protocolbuffers/protobuf: Protocol Buffers - Google's data interchange format (github.com)

    使用cmake生成vs工程直接编译即可。

    第二步：

    使用protoc命令，将编写好的proto文件生成C++类定义文件。

    .\protoc.exe --cpp_out=. ./onnx.proto

    得到 onnx.pb.h 和 onnx.pb.cpp 两个文件，将其加入到工程中。

    Onnx.proto 文件可以从 onnx/onnx: Open standard for machine learning interoperability (github.com) 找到。

    使用编译好的库文件编译的过程中，如果出现 无法解析的外部符号 "class google::protobuf::internal::ExplicitlyConstructed fixed_address_empty_string" 这个错误，添加预处理宏定义 PROTOBUF_USE_DLLS 即可。

    关键代码：

    加载反序列化onnx文件

    1.     ifstream fin("model.onnx", std::ios::in | std::ios::binary);
    2.     onnx::ModelProto onnx_model;
    3.     onnx_model.ParseFromIstream(&fin);

    打印一些信息   

    1.     std::cout << "ir_version: " << onnx_model.ir_version() << std::endl;
    2.     std::cout << "opset_import_size: " << onnx_model.opset_import_size() << std::endl;
    3.     std::cout << "OperatorSetIdProto domain: " << onnx_model.opset_import(0).domain() << std::endl;
    4.     std::cout << "OperatorSetIdProto version: " << onnx_model.opset_import(0).version() << std::endl;
    5.     std::cout << "producer_name: " << onnx_model.producer_name() << std::endl;
    6.     std::cout << "producer_version: " << onnx_model.producer_version() << std::endl;
    7.     std::cout << "domain: " << onnx_model.domain() << std::endl;
    8.     std::cout << "model_version: " << onnx_model.model_version() << std::endl;
    9.     std::cout << "doc_string: " << onnx_model.doc_string() << std::endl;

    输入节点的个数

     

    onnx_model.graph().input_size()

    输入节点的名称

     

     onnx_model.graph().input(0).name()

    输入输出节点的数据类型

    onnx_model.graph().input(0).type().tensor_type().elem_type()

    返回的是int类型，与实际数据类型对应关系见：onnx/onnx.proto at main · onnx/onnx (github.com) 文件中的 DataType 枚举类型。

    输入节点的输入维度

    1. int dim_size = onnx_model.graph().input(0).type().tensor_type().shape().dim_size();
    2. for (int i = 0; i < dim_size; i++)
    3. {
    4.     onnx_model.graph().input(0).type().tensor_type().shape().dim().Get(i).dim_value();
    5. }

    封装类 OnnxInfo

    OnnxInfo.h

    1.  
    2. #ifdef __cplusplus
    3. extern "C" {
    4. #endif
    5.  
    6.     DLL_API uint64_t onnx_load(const char* path);
    7.  
    8.     DLL_API void onnx_close(uint64_t ptr_addr);
    9.  
    10.     DLL_API int onnx_get_input_count(uint64_t ptr_addr);
    11.  
    12.     DLL_API int onnx_get_output_count(uint64_t ptr_addr);
    13.  
    14.     DLL_API const char* onnx_get_input_name(uint64_t ptr_addr, int input_index);
    15.  
    16.     DLL_API const char* onnx_get_output_name(uint64_t ptr_addr, int output_index);
    17.  
    18.     DLL_API const char* onnx_get_input_data_type(uint64_t ptr_addr, int input_index);
    19.  
    20.     DLL_API const char* onnx_get_output_data_type(uint64_t ptr_addr, int output_index);
    21.  
    22.     DLL_API int onnx_get_input_dims(uint64_t ptr_addr, int input_index, int* dims);
    23.  
    24.     DLL_API int onnx_get_output_dims(uint64_t ptr_addr, int output_index, int* dims);
    25.  
    26. #ifdef __cplusplus
    27. }
    28. #endif

    OnnxInfo.cpp

    1. #define _DLL_INTERNEL_
    2.  
    3. #include "OnnxInfo.h"
    4. #include "onnx.pb.h"
    5.  
    6. #include 
    7. #include 
    8. #include 
    9.  
    10.  
    11.  
    12. using namespace std;
    13.  
    14. #ifdef __cplusplus
    15. extern "C" {
    16. #endif
    17.  
    18.     namespace {
    19.  
    20.         char g_string[256];
    21.  
    22.         vector> g_models;
    23.  
    24.         void g_release_ptr(uint64_t ptr_addr)
    25.         {
    26.             auto iter = g_models.begin();
    27.             while (iter != g_models.end())
    28.             {
    29.                 if (ptr_addr == (uint64_t)iter->get())
    30.                 {
    31.                     g_models.erase(iter);
    32.                     break;
    33.                 }
    34.                 iter++;
    35.             }
    36.         }
    37.  
    38.         shared_ptr g_get_ptr(uint64_t ptr_addr)
    39.         {
    40.             for (auto& ptr : g_models)
    41.             {
    42.                 if (ptr_addr == (uint64_t)ptr.get())
    43.                 {
    44.                     return ptr;
    45.                 }
    46.             }
    47.             return nullptr;
    48.         }
    49.  
    50.         const char* g_get_data_type_name_by_id(int data_type_id)
    51.         {
    52.             auto dt = (onnx::TensorProto_DataType)data_type_id;
    53.             switch (dt)
    54.             {
    55.             case onnx::TensorProto_DataType_UNDEFINED:
    56.                 return "UNDEFINED";
    57.             case onnx::TensorProto_DataType_FLOAT:
    58.                 return "FLOAT";
    59.             case onnx::TensorProto_DataType_UINT8:
    60.                 return "UINT8";
    61.             case onnx::TensorProto_DataType_INT8:
    62.                 return "FLOAT";
    63.             case onnx::TensorProto_DataType_UINT16:
    64.                 return "UINT16";
    65.             case onnx::TensorProto_DataType_INT16:
    66.                 return "INT16";
    67.             case onnx::TensorProto_DataType_INT32:
    68.                 return "INT32";
    69.             case onnx::TensorProto_DataType_INT64:
    70.                 return "INT64";
    71.             case onnx::TensorProto_DataType_STRING:
    72.                 return "STRING";
    73.             case onnx::TensorProto_DataType_BOOL:
    74.                 return "BOOL";
    75.             case onnx::TensorProto_DataType_FLOAT16:
    76.                 return "FLOAT16";
    77.             case onnx::TensorProto_DataType_DOUBLE:
    78.                 return "DOUBLE";
    79.             case onnx::TensorProto_DataType_UINT32:
    80.                 return "UINT32";
    81.             case onnx::TensorProto_DataType_UINT64:
    82.                 return "UINT64";
    83.             case onnx::TensorProto_DataType_COMPLEX64:
    84.                 return "COMPLEX64";
    85.             case onnx::TensorProto_DataType_COMPLEX128:
    86.                 return "COMPLEX128";
    87.             case onnx::TensorProto_DataType_BFLOAT16:
    88.                 return "BFLOAT16";
    89.             default:
    90.                 return "";
    91.             }
    92.         }
    93.     }
    94.     
    95.     ///  加载onnx 
    96.     uint64_t onnx_load(const char* path)
    97.     {
    98.         ifstream fin(path, std::ios::in | std::ios::binary);
    99.  
    100.         auto onnx_model_ptr = make_shared();
    101.  
    102.         bool bret = onnx_model_ptr->ParseFromIstream(&fin);
    103.         fin.close();
    104.  
    105.         if (!bret) return 0;
    106.  
    107.         uint64_t ptr = (uint64_t)onnx_model_ptr.get();
    108.         g_models.push_back(move(onnx_model_ptr));
    109.         return ptr;
    110.     }
    111.  
    112.     ///  关闭onnx 
    113.     void onnx_close(uint64_t ptr_addr)
    114.     {
    115.         g_release_ptr(ptr_addr);
    116.     }
    117.  
    118.     ///  获取输入节点的个数 
    119.     int onnx_get_input_count(uint64_t ptr_addr)
    120.     {
    121.         auto ptr = g_get_ptr(ptr_addr);
    122.         if (!ptr.get()) return -1;
    123.  
    124.         if (!ptr->has_graph()) return -2;
    125.  
    126.         return ptr->graph().input_size();
    127.     }
    128.  
    129.     ///  获取输出节点的个数 
    130.     int onnx_get_output_count(uint64_t ptr_addr)
    131.     {
    132.         auto ptr = g_get_ptr(ptr_addr);
    133.         if (!ptr.get()) return -1;
    134.  
    135.         if (!ptr->has_graph()) return -2;
    136.  
    137.         return ptr->graph().output_size();
    138.     }
    139.  
    140.  
    141.     ///  获取输入节点的名称 
    142.     const char* onnx_get_input_name(uint64_t ptr_addr, int input_index)
    143.     {
    144.         auto ptr = g_get_ptr(ptr_addr);
    145.         if (!ptr.get()) return "";
    146.  
    147.         if (!ptr->has_graph()) return "";
    148.  
    149.         int input_size = ptr->graph().input_size();
    150.         if (input_index >= input_size || input_index < 0) return "";
    151.  
    152.         auto input = ptr->graph().input(input_index);
    153.  
    154.         sprintf_s(g_string, sizeof(g_string), input.name().c_str());
    155.  
    156.         return g_string;
    157.     }
    158.  
    159.     ///  获取输出节点的名称 
    160.     const char* onnx_get_output_name(uint64_t ptr_addr, int output_index)
    161.     {
    162.         auto ptr = g_get_ptr(ptr_addr);
    163.         if (!ptr.get()) return "";
    164.  
    165.         if (!ptr->has_graph()) return "";
    166.  
    167.         int output_size = ptr->graph().output_size();
    168.         if (output_index >= output_size || output_index < 0) return "";
    169.  
    170.         auto output = ptr->graph().output(output_index);
    171.  
    172.         sprintf_s(g_string, sizeof(g_string), output.name().c_str());
    173.  
    174.         return g_string;
    175.     }
    176.  
    177.     ///  获取输入节点的数据类型 
    178.     const char* onnx_get_input_data_type(uint64_t ptr_addr, int input_index)
    179.     {
    180.         auto ptr = g_get_ptr(ptr_addr);
    181.         if (!ptr.get()) return "";
    182.  
    183.         if (!ptr->has_graph()) return "";
    184.  
    185.         int input_size = ptr->graph().input_size();
    186.         if (input_index >= input_size || input_index < 0) return "";
    187.  
    188.         auto input = ptr->graph().input(input_index);
    189.  
    190.         auto type_id = input.type().tensor_type().elem_type();
    191.  
    192.         return g_get_data_type_name_by_id(type_id);
    193.     }
    194.  
    195.     ///  获取输出节点的数据类型 
    196.     const char* onnx_get_output_data_type(uint64_t ptr_addr, int output_index)
    197.     {
    198.         auto ptr = g_get_ptr(ptr_addr);
    199.         if (!ptr.get()) return "";
    200.  
    201.         if (!ptr->has_graph()) return "";
    202.  
    203.         int output_size = ptr->graph().output_size();
    204.         if (output_index >= output_size || output_index < 0) return "";
    205.  
    206.         auto output = ptr->graph().output(output_index);
    207.  
    208.         auto type_id = output.type().tensor_type().elem_type();
    209.  
    210.         return g_get_data_type_name_by_id(type_id);
    211.     }
    212.  
    213.     ///  获取输入节点的维数 
    214.     int onnx_get_input_dims(uint64_t ptr_addr, int input_index, int* dims)
    215.     {
    216.         auto ptr = g_get_ptr(ptr_addr);
    217.         if (!ptr.get()) return -1;
    218.  
    219.         if (!ptr->has_graph()) return -2;
    220.  
    221.         int input_size = ptr->graph().input_size();
    222.         if (input_index >= input_size || input_index < 0) return -3;
    223.  
    224.         auto input = ptr->graph().input(input_index);
    225.  
    226.         int dim_size = input.type().tensor_type().shape().dim_size();
    227.         
    228.  
    229.         if (dims) for (int i = 0; i < dim_size; i++)
    230.         {
    231.             dims[i] = input.type().tensor_type().shape().dim().Get(i).dim_value();
    232.         }
    233.  
    234.         return dim_size;
    235.     }
    236.  
    237.     ///  获取输出节点的维数 
    238.     int onnx_get_output_dims(uint64_t ptr_addr, int output_index, int* dims)
    239.     {
    240.         auto ptr = g_get_ptr(ptr_addr);
    241.         if (!ptr.get()) return -1;
    242.  
    243.         if (!ptr->has_graph()) return -2;
    244.  
    245.         int output_size = ptr->graph().output_size();
    246.         if (output_index >= output_size || output_index < 0) return -3;
    247.  
    248.         auto output = ptr->graph().output(output_index);
    249.  
    250.         int dim_size = output.type().tensor_type().shape().dim_size();
    251.  
    252.         if (dims) for (int i = 0; i < dim_size; i++)
    253.         {
    254.             dims[i] = output.type().tensor_type().shape().dim().Get(i).dim_value();
    255.         }
    256.  
    257.         return dim_size;
    258.     }
    259.  
    260.  
    261.  
    262. #ifdef __cplusplus
    263. }
    264. #endif

    调用方法

    1. int main()
    2. {
    3.     const char* path = "myunet.onnx";
    4.  
    5.     auto adr = onnx_load(path);
    6.  
    7.     int ic = onnx_get_input_count(adr);
    8.     int oc = onnx_get_output_count(adr);
    9.  
    10.     cout << onnx_get_input_name(adr, 0) << endl;
    11.     cout << onnx_get_output_name(adr, 0) << endl;
    12.  
    13.     cout << onnx_get_input_data_type(adr, 0) << endl;
    14.     cout << onnx_get_output_data_type(adr, 0) << endl;
    15.  
    16.     int dim_size;
    17.     int dims[8];
    18.  
    19.     cout << "input dim size: " << onnx_get_input_dims(adr, 0, 0) << endl;
    20.     cout << "input dim size2: " << (dim_size = onnx_get_input_dims(adr, 0, dims)) << endl;
    21.     for (int i = 0; i < dim_size; i++)
    22.     {
    23.         cout << dims[i] << " ";
    24.     }
    25.     cout << endl;
    26.  
    27.     cout << "output dim size: " << onnx_get_output_dims(adr, 0, 0) << endl;
    28.     cout << "output dim size2: " << (dim_size = onnx_get_output_dims(adr, 0, dims)) << endl;
    29.     for (int i = 0; i < dim_size; i++)
    30.     {
    31.         cout << dims[i] << " ";
    32.     }
    33.     cout << endl;
    34.  
    35.     onnx_close(adr);
    36.  
    37.     cin.ignore();
    38.  
    39.     return 0;
    40. }

  • 相关阅读:
    Kafka干货之「零拷贝」
    基于JSP的九宫格日志网站
    Python与数据分析--Pandas-1
    计算机系统的层次结构
    tcpdump抓包详解
    【API篇】十一、Flink水位线传递与迟到数据处理
    STM32驱动AHT10&OLED显示温湿度
    FFmpeg开发笔记（九）Linux交叉编译Android的x265库
    Chapter9 : De Novo Molecular Design with Chemical Language Models
    第十七节 huggingface的trainner的断点续训的Demo(resume)
  • 原文地址：https://blog.csdn.net/Ango_/article/details/127790921