mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-14 20:48:00 +00:00
80 lines
3.8 KiB
Python
80 lines
3.8 KiB
Python
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
|
|
import argparse
|
|
import onnx
|
|
import os
|
|
import sys
|
|
|
|
from onnx.onnx_pb import ModelProto
|
|
from google.protobuf import text_format
|
|
|
|
onnx_operator_list = [
|
|
"Abs", "Acos", "Add", "And", "ArgMax", "ArgMin", "Asin", "Atan", "AveragePool", "BatchNormalization",
|
|
"Cast", "Ceil", "Clip", "Concat", "Constant", "Conv", "ConvTranspose", "Cos", "DepthToSpace", "Div",
|
|
"Dropout", "Elu", "Equal", "Exp", "Flatten", "Floor", "GRU", "Gather", "Gemm", "GlobalAveragePool",
|
|
"GlobalLpPool", "GlobalMaxPool", "Greater", "HardSigmoid", "Hardmax", "Identity",
|
|
"InstanceNormalization", "LRN", "LSTM", "LeakyRelu", "Less", "Log", "LogSoftmax", "LpNormalization",
|
|
"LpPool", "MatMul", "Max", "MaxPool", "MaxRoiPool", "Mean", "Min", "Mul", "Multinomial", "Neg", "Not",
|
|
"Or", "PRelu", "Pad", "Pow", "RNN", "RandomNormal", "RandomNormalLike", "RandomUniform", "RandomUniformLike",
|
|
"Reciprocal", "ReduceL1", "ReduceL2", "ReduceLogSum", "ReduceLogSumExp", "ReduceMax", "ReduceMean",
|
|
"ReduceMin", "ReduceProd", "ReduceSum", "ReduceSumSquare", "Relu", "Reshape", "Selu", "Shape", "Sigmoid",
|
|
"Sin", "Size", "Slice", "Softmax", "Softplus", "Softsign", "SpaceToDepth", "Split", "Sqrt", "Squeeze",
|
|
"Sub", "Sum", "Tan", "Tanh", "Tile", "TopK", "Transpose", "Unsqueeze", "Upsample", "Xor",
|
|
"ATen", "Affine", "ConstantFill", "Crop", "GRUUnit", "GivenTensorFill", "If", "ImageScaler", "Loop",
|
|
"LoopIndexTensor", "MeanVarianceNormalization", "ParametricSoftplus", "Scale", "ScaledTanh", "ThresholdedRelu"
|
|
]
|
|
|
|
#Given a model directory, it parses the model operators and outputs in CSV formst as below:
|
|
# Abs, Acos, Add, Add, ...
|
|
# Model1, 5, 18, , 1, ...
|
|
# Model2, , 19, 3, , ...
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description='Operators of ONNX model.')
|
|
parser.add_argument('modelDir', nargs=1, help='Model Directory Path')
|
|
parser.add_argument('resultFile', nargs=1, help='Result File Path')
|
|
|
|
model_dir = parser.parse_args().modelDir[0]
|
|
result_file = parser.parse_args().resultFile[0]
|
|
|
|
with open(result_file, 'a') as result:
|
|
|
|
for root, dirs, files in os.walk(model_dir):
|
|
for file in files:
|
|
if file.endswith('.onnx'):
|
|
abs_path = root + "/" + file
|
|
f = open(abs_path, "rb")
|
|
proto = ModelProto()
|
|
proto.ParseFromString(f.read())
|
|
f.close()
|
|
|
|
# Get model operator list
|
|
operators = {}
|
|
nodes = proto.graph.node
|
|
for node in nodes:
|
|
print(node.op_type)
|
|
if not node.op_type in operators:
|
|
operators[node.op_type]=1
|
|
else:
|
|
operators[node.op_type]+=1
|
|
|
|
# store all the operators with format:
|
|
# ModelName[,occurences]*
|
|
|
|
#head
|
|
head = ''
|
|
for op in onnx_operator_list:
|
|
head += ',{}'.format(op)
|
|
result.write('{}\n'.format(head))
|
|
|
|
#content
|
|
model_name = file[:-5]
|
|
if model_name == 'model':
|
|
model_name = os.path.basename(os.path.dirname(abs_path)) # use directory name
|
|
|
|
model_operators_string = '{}'.format(model_name)
|
|
for op in onnx_operator_list:
|
|
model_operators_string+=','
|
|
if op in operators:
|
|
model_operators_string+='{}'.format(operators[op])
|
|
result.write('{}\n'.format(model_operators_string))
|