如何找到ONNX模型的输入大小?我最终想从python脚本。
使用TensorFlow i可以恢复图形定义,从中找到输入候选节点,然后获取它们的大小。我可以使用onnx(甚至更简单(做类似的事情吗?
谢谢
是的,前提是输入模型具有信息。请注意,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()
请不要将input
用作变量名称,因为它是内置功能。
想到的第一个想法是,如果我需要姓名,data_type或Protobuf对象的某些属性,请使用google.protobuf.json_format.MessageToDict()
方法。例如:
form google.protobuf.json_format import MessageToDict
model = onnx.load("path/to/model.onnx")
for _input in model.graph.input:
print(MessageToDict(_input))
将给出类似的输出:
{'name': '0', 'type': {'tensorType': {'elemType': 2, 'shape': {'dim': [{'dimValue': '4'}, {'dimValue': '3'}, {'dimValue': '384'}, {'dimValue': '640'}]}}}}
我不太清楚每个model.graph.input
是否是RepeatedCompositeContainer
对象,但是当它是RepeatedCompositeContainer
时,必须使用for
循环。
然后您需要从dim
字段获取形状信息。
model = onnx.load("path/to/model.onnx")
for _input in model.graph.input:
m_dict = MessageToDict(_input))
dim_info = m_dict.get("type").get("tensorType").get("shape").get("dim") # ugly but we have to live with this when using dict
input_shape = [d.get("dimValue") for d in dim_info] # [4,3,384,640]
如果您唯一需要DIM,请改用消息对象。
model = onnx.load("path/to/model.onnx")
for _input in model.graph.input:
dim = _input.type.tensor_ype.shape.dim
input_shape = [MessgeToDict(d).get("dimValue") for d in dim]
# if you prefer the python naming style, using the line below
# input_shape = [MessgeToDict(d, preserving_proto_field_name=True).get("dim_value") for d in dim]
一行版本:
model = onnx.load("path/to/model.onnx")
input_shapes = [[d.dim_value for d in _input.type.tensor_type.shape.dim] for _input in model.graph.input]
参考:
https://github.com/googleapis/python-vision/issues/70
attributeError:'google.protobuf.pyext._message._message.repeatedcompositeco'对象没有属性'append'
如果您使用onnxRuntime而不是onnx进行推理。
尝试使用以下代码。
import onnxruntime as ort
model = ort.InferenceSession("model.onnx", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
input_shape = model.get_inputs()[0].shape