前言

分析每一層 onnx model 的內容

範例程式碼

import onnx

model_path = "model.onnx"
model = onnx.load(model_path)

inferred_model = onnx.shape_inference.infer_shapes(model)

# input
print("> Model Inputs:")
for input in inferred_model.graph.input:
    input_shape = [dim.dim_value for dim in input.type.tensor_type.shape.dim]
    print(f"Input: {input.name}, Shape: {input_shape}")

# output edge name
print("> Hidden Layers:")
# for node in inferred_model.graph.value_info:
#     shape = [dim.dim_value for dim in node.type.tensor_type.shape.dim]
#     print(f"Node: {node.name}, Shape: {shape}")


# output layer name with I/O shape
def get_node_shapes(node_outputs):
    shapes = []
    for output in node_outputs:
        for value_info in inferred_model.graph.value_info:
            if value_info.name == output:
                shape = [dim.dim_value for dim in value_info.type.tensor_type.shape.dim]
                shapes.append(shape)
    return shapes


for node in inferred_model.graph.node:
    node_type = node.op_type
    node_input_shapes = get_node_shapes(node.input)
    node_output_shapes = get_node_shapes(node.output)

    print(
        f"Node Name: {node.name}, Type: {node_type}, Input Shapes: {node_input_shapes}, Output Shapes: {node_output_shapes}"
    )

# output
print("> Model Outputs:")
for output in inferred_model.graph.output:
    output_shape = [dim.dim_value for dim in output.type.tensor_type.shape.dim]
    print(f"Output: {output.name}, Shape: {output_shape}")