python - 从 onnx 文件中查找输入形状

Question

如何找到 onnx 模型的输入大小？我最终想从 python 编写脚本。

使用 tensorflow，我可以恢复图形定义，从中找到输入候选节点，然后获取它们的大小。我可以用 ONNX（甚至更简单）做类似的事情吗？

谢谢

标签： pythononnx

Answer 1

是的，只要输入模型有信息。请注意，ONNX 模型的输入可能具有未知等级，或者具有固定（如 100）或符号（如“N”）或完全未知的维度的已知等级。您可以按如下方式访问：

import onnx

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

# The model is represented as a protobuf structure and it can be accessed
# using the standard python-for-protobuf methods

# iterate through inputs of the graph
for input in model.graph.input:
    print (input.name, end=": ")
    # get type of input tensor
    tensor_type = input.type.tensor_type
    # check if it has a shape:
    if (tensor_type.HasField("shape")):
        # iterate through dimensions of the shape:
        for d in tensor_type.shape.dim:
            # the dimension may have a definite (integer) value or a symbolic identifier or neither:
            if (d.HasField("dim_value")):
                print (d.dim_value, end=", ")  # known dimension
            elif (d.HasField("dim_param")):
                print (d.dim_param, end=", ")  # unknown dimension with symbolic name
            else:
                print ("?", end=", ")  # unknown dimension with no name
    else:
        print ("unknown rank", end="")
    print()