ONNX 解析打印 initializer

import onnx
import sys
import numpy as np

def onnx_datatype_to_npType(data_type):
    if data_type == 1:
        return np.float32
    else:
        raise TypeError("don't support data type")
    
if __name__ == "__main__":
    
    model_name = "test.onnx"
    
    onnx_model = onnx.load(model_name)
    
    graph = onnx_model.graph
    nodes = graph.node
    initializer = graph.initializer
    inputs= graph.input
    outputs = graph.output

    name_lists = ["Upsample_260","Upsample_267","Upsample_274","Upsample_281"]
    
    for i in range(len(initializer)):
        if initializer[i].name in name_lists:
            print(initializer[i].name,"\t",end="")
            dtype = initializer[i].data_type
            params = np.frombuffer(initializer[i].raw_data, dtype=onnx_datatype_to_npType(dtype))
            print(params,end="\n")

你可能感兴趣的:(ONNX,python,机器学习,人工智能)