onnx模型数据结构分析,用于解析onnx模型

文章目录onnx 模型数据结构一、核心字段分类二、关键嵌套结构详解1. GraphProto计算图核心node2. OperatorSetProto算子集依赖3. TensorProto张量定义onnx存储为json结构onnx 模型数据结构onnx是基于protobuffer的, scheme 定义在:https://github.com/onnx/onnx/blob/main/onnx/onnx.proto3onnx官方其实也有一篇文章在介绍onnx IR的结构: onnx IR一、核心字段分类ModelProto 的字段可分为基础标识 / 版本、元数据、计算图核心、扩展能力四大类以下是关键字段的详细说明可以看到整个模型是ModelProto:字段类别核心字段作用说明版本与兼容性ir_version标识 ONNX 中间表示IR的版本如 8/9/13 等决定语法 / 语义兼容性opset_import模型依赖的算子集OperatorSet列表格式为{domain: version}如 : 19表示默认域的 19 版算子元数据producer_name生成模型的工具 / 框架名称如 PyTorch/TensorFlowproducer_version生成工具的版本号model_version模型自身的版本自定义整数用于版本管理doc_string模型的可读文档支持 Markdownmetadata_props自定义键值对元数据如作者、训练数据集、license 等计算图核心graph模型的计算逻辑核心类型为GraphProto见下文拆解扩展能力functions模型内自定义函数FunctionProto用于封装可复用的算子逻辑training_info训练相关信息TrainingInfoProto存储训练步骤、优化器等训练态信息二、关键嵌套结构详解ModelProto 依赖多个子 Proto 结构实现完整能力核心嵌套结构如下1. GraphProto计算图核心graph是 ModelProto 的核心定义了模型的计算逻辑其核心字段name图的名称唯一标识input/output图的输入 / 输出张量定义ValueInfoProto包含张量名称、类型TypeProto、形状TensorShapeProtonode计算节点列表NodeProto每个节点对应一个算子如 Conv/Add/MatMul包含op_type算子类型、input/output张量名、attribute算子属性如卷积核大小initializer模型的权重 / 常量张量TensorProto如卷积核、偏置等sparse_initializer稀疏张量类型的权重可选。字段名数据类型核心作用补充说明namestring计算图的唯一标识名称如cspresnet50_graph用于区分不同子图如循环 / 条件子图inputrepeated ValueInfoProto图的输入张量元信息不含权重每个ValueInfoProto包含张量名、类型TypeProto、形状TensorShapeProto如模型输入input.1outputrepeated ValueInfoProto图的输出张量元信息如模型输出output.1定义输出张量的名称、类型、形状noderepeated NodeProto计算图的核心算子节点DAG 节点每个NodeProto对应一个算子如 Conv/Relu包含输入、输出、属性AttributeProtoinitializerrepeated TensorProto模型权重 / 常量张量如卷积核、偏置存储张量的二进制数据raw_data、形状dims、数据类型data_type是模型参数的核心载体sparse_initializerrepeated SparseTensorProto稀疏权重张量如稀疏卷积的权重仅用于稀疏张量场景包含非零值values、索引indices等doc_stringstring计算图的文档说明备注 / 注释可选字段用于描述图的用途、版本等value_inforepeated ValueInfoProto图中中间张量的元信息非输入输出 / 初始化器可选字段用于补充中间张量如 Conv 输出的类型 / 形状信息quantization_annotationrepeated QuantizationAnnotationProto量化相关标注如张量的量化参数量化模型专用描述张量的零点、缩放因子等layout_annotationrepeated LayoutAnnotationProto张量布局标注如 NCHW/NHWC可选字段用于指定张量的维度顺序metadata_propsrepeated StringStringEntryProto自定义元数据键值对如author: alan、version: 1.0functionsrepeated FunctionProto图内自定义算子函数用于封装复用的算子逻辑如自定义 Selu 函数argrepeated AttributeProto图级别的属性极少使用区别于NodeProto的算子属性用于全局配置nodeprotomessage NodeProto { string name 1; // 节点名可选 string op_type 2; // 算子类型如 Conv/Relu/MatMul repeated string input 3;// 输入张量名列表如 [input.1, conv1.weight] repeated string output 4;// 输出张量名列表如 [conv1.output] string domain 5; // 算子域默认 为 ONNX 标准算子如 com.microsoft 为自定义算子 repeated AttributeProto attribute 6; // 算子属性如 Conv 的 stride/padding string doc_string 10; // 节点说明 }json{input:[/stem/conv1/conv/Conv_output_0],output:[/stem/conv1/bn/act/LeakyRelu_output_0],name:/stem/conv1/bn/act/LeakyRelu,op_type:LeakyRelu,attribute:[{name:alpha,f:0.01,type:FLOAT}]},2. OperatorSetProto算子集依赖opset_import引用的算子集定义核心字段domain算子集的域名默认域为空字符串自定义算子可指定如custom_domainversion算子集版本如 15/19决定算子的签名 / 语义。3. TensorProto张量定义用于描述权重、输入输出张量核心字段data_type张量元素类型如FLOAT/INT32/STRING对应枚举值 1/6/7dims张量形状如[3, 224, 224]表示 3 通道 224x224 图像raw_data张量的二进制数据序列化后的权重值辅助字段name张量名、segment稀疏张量分段信息等。onnx存储为json结构importonnximportjsonfromgoogle.protobuf.json_formatimportMessageToDict model:onnx.ModelProtoonnx.load(/Users/alanchen/workspace/cspresnet50_Opset16.onnx)graph:onnx.GraphProtomodel.graphdefstrip_tensor_weights(graph_dict:dict)-dict: 剥离 Graph 字典中 Tensor 的权重数据仅保留元信息名称、形状、数据类型 # 处理 initializer模型权重/常量张量ifinitializeringraph_dict:fortensoringraph_dict[initializer]:# 移除存储权重的字段按需扩展覆盖所有数值存储字段weight_fields[raw_data,# 二进制权重最常见float_data,# float 类型数值列表int32_data,# int32 类型数值列表int64_data,# int64 类型数值列表uint64_data,# uint64 类型数值列表bool_data,# bool 类型数值列表string_data,# string 类型数值列表double_data,# double 类型数值列表int8_data,# int8 类型数值列表uint8_data,# uint8 类型数值列表int16_data,# int16 类型数值列表uint16_data,# uint16 类型数值列表]forfieldinweight_fields:iffieldintensor:deltensor[field]# 可选处理 sparse_initializer稀疏张量权重若有ifsparse_initializeringraph_dict:forsparse_tensoringraph_dict[sparse_initializer]:ifvaluesinsparse_tensor:delsparse_tensor[values]# 移除稀疏张量的数值ifindicesinsparse_tensor:delsparse_tensor[indices]# 移除稀疏张量的索引returngraph_dict# 步骤1将 GraphProto 转为字典Protobuf 原生转换graph_dictMessageToDict(graph,preserving_proto_field_nameTrue)# 步骤2剥离权重数据graph_dict_strippedstrip_tensor_weights(graph_dict)# 步骤3导出为 JSON 文件无权重withopen(cspresnet50_graph_proto.json,w,encodingutf-8)asf:json.dump(graph_dict_stripped,f,ensure_asciiFalse,indent2,sort_keysFalse)print(已导出剥离权重的 Graph JSON 文件cspresnet50_graph_proto_stripped.json)