python简单解析打印onnx模型信息

当我们加载了一个ONNX之后，我们获得的就是一个ModelProto，它包含了一些版本信息，生产者信息和一个GraphProto。在GraphProto里面又包含了四个repeated数组，它们分别是node(NodeProto类型)，input(ValueInfoProto类型)，output(ValueInfoProto类型)和initializer(TensorProto类型)，其中node中存放了模型中所有的计算节点，input存放了模型的输入节点，output存放了模型中所有的输出节点，initializer存放了模型的所有权重参数。

这里用python，将onnx中包含的有用的信息打印出来，进行一个直观可视化。

整体流程：

• 解析graph input
• 解析graph output
• 解析graph initializer，模型的所有权重参数
• 解析graph node，打印op信息，输入输出，得到整个计算的graph

import onnx
import numpy as np
import logging
logging.basicConfig(level=logging.INFO)

def onnx_datatype_to_npType(data_type):
    if data_type == 1:
        return np.float32
    else:
        raise TypeError("don't support data type")


def parser_initializer(initializer):
    name = initializer.name
    logging.info(f"initializer name: {name}")

    dims = initializer.dims
    shape = [x for x in dims]
    logging.info(f"initializer with shape:{shape}")

    dtype = initializer.data_type
    logging.info(f"initializer with type: {onnx_datatype_to_npType(dtype)} ")
    
    # print tenth buffer
    weights = np.frombuffer(initializer.raw_data, dtype=onnx_datatype_to_npType(dtype))
    logging.info(f"initializer first 10 wights:{weights[:10]}")



def parser_tensor(tensor, use='normal'):
    name = tensor.name
    logging.info(f"{use} tensor name: {name}")

    data_type = tensor.type.tensor_type.elem_type
    logging.info(f"{use} tensor data type: {data_type}")

    dims = tensor.type.tensor_type.shape.dim
    shape = []
    for i, dim in enumerate(dims):
        shape.append(dim.dim_value)
    logging.info(f"{use} tensor with shape:{shape} ")


def parser_node(node):
    def attri_value(attri):
        if attri.type == 1:
            return attri.i
        elif attri.type == 7:
            return list(attri.ints)
        
    name = node.name
    logging.info(f"node name:{name}")

    opType = node.op_type
    logging.info(f"node op type:{opType}")

    inputs = list(node.input)
    logging.info(f"node with {len(inputs)} inputs:{inputs}")

    outputs = list(node.output)
    logging.info(f"node with {len(outputs)} outputs:{outputs}")
    
    attributes = node.attribute
    for attri in attributes:
        name = attri.name
        value = attri_value(attri)
        logging.info(f"{name} with value:{value}")


def parser_info(onnx_model):
    ir_version = onnx_model.ir_version
    producer_name = onnx_model.producer_name
    producer_version = onnx_model.producer_version
    for info in [ir_version, producer_name, producer_version]:
        logging.info("onnx model with info:{}".format(info))

def parser_inputs(onnx_graph):
    inputs = onnx_graph.input
    for input in inputs:
        parser_tensor(input, 'input')

def parser_outputs(onnx_graph):
    outputs = onnx_graph.output
    for output in outputs:
        parser_tensor(output, 'output')

def parser_graph_initializers(onnx_graph):
    initializers = onnx_graph.initializer
    for initializer in initializers:
        parser_initializer(initializer)


def parser_graph_nodes(onnx_graph):
    nodes = onnx_graph.node
    for node in nodes:
        parser_node(node)
        t = 1

def onnx_parser():
    model_path = 'D:/project/public/yolov5-5.0/yolov5s-sim.onnx'
    model = onnx.load(model_path)

    # 0.
    parser_info(model)

    graph = model.graph

    # 1.
    parser_inputs(graph)

    # 2. 
    parser_outputs(graph)

    # 3.
    parser_graph_initializers(graph)

    # 4. 
    parser_graph_nodes(graph)


if __name__ == '__main__':
    onnx_parser()

INFO:root:onnx model with info:7
INFO:root:onnx model with info:pytorch
INFO:root:onnx model with info:1.10
INFO:root:input tensor name: images
INFO:root:input tensor data type: 1
INFO:root:input tensor with shape:[1, 3, 640, 640]
INFO:root:output tensor name: output
INFO:root:output tensor data type: 1
INFO:root:output tensor with shape:[1, 3, 80, 80, 85]
INFO:root:output tensor name: 405
INFO:root:output tensor data type: 1
INFO:root:output tensor with shape:[1, 3, 40, 40, 85]
INFO:root:output tensor name: 419
INFO:root:output tensor data type: 1
INFO:root:output tensor with shape:[1, 3, 20, 20, 85]
INFO:root:initializer name: model.0.conv.conv.weight
INFO:root:initializer with shape:[32, 12, 3, 3]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[-0.2730072  -1.4410826  -1.187087   -0.31312177 -0.94754034 -0.7239634    
  0.4315643   2.0547783   1.5080036   0.04422583]
INFO:root:initializer name: model.0.conv.conv.bias
INFO:root:initializer with shape:[32]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[ 1.4819448  -0.9827409  -0.7623507  -0.503065   -0.1677513   3.1156974    
  1.4849529   0.15412715 -0.4954783   2.8073668 ]
INFO:root:initializer name: model.1.conv.weight
INFO:root:initializer with shape:[64, 32, 3, 3]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[-0.0130239  -0.00788062  0.0033419  -0.08140857 -0.1592819  -0.10095055   
 -0.01685373 -0.00378994 -0.01589627  0.01994706]
INFO:root:initializer name: model.1.conv.bias
INFO:root:initializer with shape:[64]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[ 2.9810083   1.6583345   2.4536395   4.068509   -0.03429741 -0.2963567    
  1.7936802   0.32334638  2.3112206   0.9088864 ]
INFO:root:initializer name: model.2.cv1.conv.weight
INFO:root:initializer with shape:[32, 64, 1, 1]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[ 0.00402472 -0.02254777  0.01670638  0.07000323 -0.01463853 -0.01208      
 -0.00047498 -0.01327164  0.02764146  0.03544556]
...
...
...
INFO:root:initializer name: model.23.cv2.conv.bias
INFO:root:initializer with shape:[256]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[ 0.33130765  0.50842535  0.04622685  0.27884385  0.35381323 -0.05325952   
  0.18430158  0.16491716 -0.4574555   0.21860392]
INFO:root:initializer name: model.23.cv3.conv.weight
INFO:root:initializer with shape:[512, 512, 1, 1]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[ 0.15170689 -0.05846716 -0.03708049 -0.00515915  0.04973555 -0.01793442   
  0.26971823 -0.08900855 -0.08623905 -0.01718434]
INFO:root:initializer name: model.23.cv3.conv.bias
INFO:root:initializer with shape:[512]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[-0.07636432 -0.14711903 -0.1693097   1.2009852   0.00255883  1.4637253
 -0.29597348  2.088275    0.5806205   0.49393153]
INFO:root:initializer name: model.23.m.0.cv1.conv.weight
INFO:root:initializer with shape:[256, 256, 1, 1]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[ 1.8370461e-01 -3.6310073e-02  2.6413003e-02 -3.4686580e-02
 -4.4441203e-04  3.3812389e-02 -3.7558913e-02 -9.6223257e-02
 -5.3294320e-02 -5.8845425e-01]
INFO:root:initializer name: model.23.m.0.cv1.conv.bias
INFO:root:initializer with shape:[256]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[ 0.4050864  -0.03902826  0.31003702  
INFO:root:initializer name: model.24.m.0.bias
INFO:root:initializer with shape:[255]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[ 0.13708496  0.06561279  0.81152344  0.62353516 -4.1992188  -1.0546875    
 -5.2734375  -3.4414062  -5.5234375  -5.71875   ]
INFO:root:initializer name: model.24.m.1.weight
INFO:root:initializer with shape:[255, 256, 1, 1]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[-0.00010455 -0.00139713  0.00134754  0.01174927 -0.00017703 -0.00750351
 -0.00029159 -0.00182247 -0.02702332  0.04980469]
INFO:root:initializer name: model.24.m.1.bias
INFO:root:initializer with shape:[255]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[-0.05148315 -0.05493164  0.6333008   0.05197144 -2.625      -1.3398438    
 -5.7851562  -4.765625   -5.7929688  -6.2773438 ]
INFO:root:initializer name: model.24.m.2.weight
INFO:root:initializer with shape:[255, 512, 1, 1]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[-2.5444031e-03 -1.1138916e-01  9.1195107e-06  1.0973215e-04
 -1.5907288e-03 -3.3130646e-03  2.6941299e-04 -9.5486641e-05
  1.4615059e-04 -7.8857422e-02]
INFO:root:initializer name: model.24.m.2.bias
INFO:root:initializer with shape:[255]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[-1.6265869e-02 -1.9702911e-03 -6.9091797e-02  1.9494629e-01
 -2.0878906e+00 -1.6210938e+00 -6.5625000e+00 -5.4531250e+00
 -6.3437500e+00 -6.7070312e+00]
INFO:root:initializer name: 420
INFO:root:initializer with shape:[4]
INFO:root:initializer with type: <class 'numpy.float32'>
INFO:root:initializer first 10 wights:[1. 1. 2. 2.]