wehub-resource-sync
337 lines
13 KiB
Python
337 lines
13 KiB
Python
""" ONNX metadata script """
|
|
|
|
import collections
|
|
import json
|
|
import os
|
|
import re
|
|
|
|
import onnx.backend.test.case
|
|
import onnx.defs
|
|
import onnx.onnx_ml_pb2
|
|
import onnxruntime
|
|
|
|
attribute_type_table = [
|
|
"undefined",
|
|
"float32",
|
|
"int64",
|
|
"string",
|
|
"tensor",
|
|
"graph",
|
|
"float32[]",
|
|
"int64[]",
|
|
"string[]",
|
|
"tensor[]",
|
|
"graph[]",
|
|
"sparse_tensor",
|
|
"sparse_tensor[]",
|
|
"type_proto",
|
|
"type_proto[]"
|
|
]
|
|
|
|
|
|
def _format_description(description):
|
|
def replace_line(match):
|
|
link = match.group(1)
|
|
url = match.group(2)
|
|
if not url.startswith("http://") and not url.startswith("https://"):
|
|
url = "https://github.com/onnx/onnx/blob/master/docs/" + url
|
|
return "[" + link + "](" + url + ")"
|
|
return re.sub('\\[(.+)\\]\\(([^ ]+?)( "(.+)")?\\)', replace_line, description)
|
|
|
|
def _format_range(value):
|
|
return "∞" if value == 2147483647 else str(value)
|
|
|
|
class OnnxSchema:
|
|
""" ONNX schema """
|
|
|
|
def __init__(self, schema, snippets):
|
|
self.schema = schema
|
|
self.snippets = snippets
|
|
self.name = self.schema.name
|
|
self.module = self.schema.domain if self.schema.domain else "ai.onnx"
|
|
self.version = self.schema.since_version
|
|
self.key = self.name + ":" + self.module + ":" + str(self.version).zfill(4)
|
|
|
|
def _get_attr_type(self, attribute_type, attribute_name, op_type, op_domain):
|
|
key = op_domain + ":" + op_type + ":" + attribute_name
|
|
if key in (":Cast:to", ":EyeLike:dtype", ":RandomNormal:dtype"):
|
|
return "DataType"
|
|
return attribute_type_table[attribute_type]
|
|
|
|
def _get_attr_default_value(self, attr_value):
|
|
if attr_value.HasField("i"):
|
|
return attr_value.i
|
|
if attr_value.HasField("s"):
|
|
return attr_value.s.decode("utf8")
|
|
if attr_value.HasField("f"):
|
|
return attr_value.f
|
|
return None
|
|
|
|
def _update_attributes(self, value, schema):
|
|
target = value["attributes"] = []
|
|
attributes = sorted(schema.attributes.items())
|
|
for _ in collections.OrderedDict(attributes).values():
|
|
value = {}
|
|
value["name"] = _.name
|
|
attr_type = self._get_attr_type(_.type, _.name, schema.name, schema.domain)
|
|
if attr_type:
|
|
value["type"] = attr_type
|
|
value["required"] = _.required
|
|
default_value = self._get_attr_default_value(_.default_value)
|
|
if default_value:
|
|
value["default"] = default_value
|
|
description = _format_description(_.description)
|
|
if len(description) > 0:
|
|
value["description"] = description
|
|
target.append(value)
|
|
|
|
def _update_inputs(self, value, inputs):
|
|
target = value["inputs"] = []
|
|
for _ in inputs:
|
|
value = {}
|
|
value["name"] = _.name
|
|
value["type"] = _.type_str
|
|
if _.option == onnx.defs.OpSchema.FormalParameterOption.Optional:
|
|
value["option"] = "optional"
|
|
elif _.option == onnx.defs.OpSchema.FormalParameterOption.Variadic:
|
|
value["list"] = True
|
|
description = _format_description(_.description)
|
|
if len(description) > 0:
|
|
value["description"] = description
|
|
target.append(value)
|
|
|
|
def _update_outputs(self, value, outputs):
|
|
target = value["outputs"] = []
|
|
for _ in outputs:
|
|
value = {}
|
|
value["name"] = _.name
|
|
value["type"] = _.type_str
|
|
if _.option == onnx.defs.OpSchema.FormalParameterOption.Optional:
|
|
value["option"] = "optional"
|
|
elif _.option == onnx.defs.OpSchema.FormalParameterOption.Variadic:
|
|
value["list"] = True
|
|
description = _format_description(_.description)
|
|
if len(description) > 0:
|
|
value["description"] = description
|
|
target.append(value)
|
|
|
|
def _update_type_constraints(self, value, type_constraints):
|
|
value["type_constraints"] = []
|
|
for _ in type_constraints:
|
|
value["type_constraints"].append({
|
|
"description": _.description,
|
|
"type_param_str": _.type_param_str,
|
|
"allowed_type_strs": _.allowed_type_strs
|
|
})
|
|
|
|
def _update_snippets(self, value, snippets):
|
|
target = value["examples"] = []
|
|
for summary, code in sorted(snippets):
|
|
lines = code.splitlines()
|
|
while len(lines) > 0 and re.search("\\s*#", lines[-1]):
|
|
lines.pop()
|
|
if len(lines) > 0 and len(lines[-1]) == 0:
|
|
lines.pop()
|
|
target.append({
|
|
"summary": summary,
|
|
"code": "\n".join(lines)
|
|
})
|
|
|
|
def to_dict(self):
|
|
""" Serialize model to JSON message """
|
|
value = {}
|
|
value["name"] = self.name
|
|
value["module"] = self.module
|
|
value["version"] = self.version
|
|
if self.schema.support_level != onnx.defs.OpSchema.SupportType.COMMON:
|
|
value["status"] = self.schema.support_level.name.lower()
|
|
description = _format_description(self.schema.doc.lstrip())
|
|
if len(description) > 0:
|
|
value["description"] = description
|
|
if self.schema.attributes:
|
|
self._update_attributes(value, self.schema)
|
|
if self.schema.inputs:
|
|
self._update_inputs(value, self.schema.inputs)
|
|
value["min_input"] = self.schema.min_input
|
|
value["max_input"] = self.schema.max_input
|
|
if self.schema.outputs:
|
|
self._update_outputs(value, self.schema.outputs)
|
|
value["min_output"] = self.schema.min_output
|
|
value["max_output"] = self.schema.max_output
|
|
if self.schema.min_input != self.schema.max_input:
|
|
value["inputs_range"] = _format_range(self.schema.min_input) + " - " \
|
|
+ _format_range(self.schema.max_input)
|
|
if self.schema.min_output != self.schema.max_output:
|
|
value["outputs_range"] = _format_range(self.schema.min_output) + " - " \
|
|
+ _format_range(self.schema.max_output)
|
|
if self.schema.type_constraints:
|
|
self._update_type_constraints(value, self.schema.type_constraints)
|
|
if self.name in self.snippets:
|
|
self._update_snippets(value, self.snippets[self.name])
|
|
return value
|
|
|
|
class OnnxRuntimeSchema:
|
|
""" ONNX Runtime schema """
|
|
|
|
def __init__(self, schema):
|
|
self.schema = schema
|
|
self.name = self.schema.name
|
|
self.module = self.schema.domain if self.schema.domain else "ai.onnx"
|
|
self.version = self.schema.since_version
|
|
self.key = self.name + ":" + self.module + ":" + str(self.version).zfill(4)
|
|
|
|
def _get_attr_type(self, attribute_type):
|
|
return attribute_type_table[attribute_type]
|
|
|
|
def _get_attr_default_value(self, attr_value):
|
|
if attr_value.HasField("i"):
|
|
return attr_value.i
|
|
if attr_value.HasField("s"):
|
|
return attr_value.s.decode("utf8")
|
|
if attr_value.HasField("f"):
|
|
return attr_value.f
|
|
return None
|
|
|
|
def _update_attributes(self, value, schema):
|
|
target = value["attributes"] = []
|
|
attributes = sorted(schema.attributes.items())
|
|
for _ in collections.OrderedDict(attributes).values():
|
|
value = {}
|
|
value["name"] = _.name
|
|
attribute_type = self._get_attr_type(_.type)
|
|
if attribute_type:
|
|
value["type"] = attribute_type
|
|
value["required"] = _.required
|
|
default_value = onnx.onnx_ml_pb2.AttributeProto()
|
|
default_value.ParseFromString(_._default_value)
|
|
default_value = self._get_attr_default_value(default_value)
|
|
if default_value:
|
|
value["default"] = default_value
|
|
description = _format_description(_.description)
|
|
if len(description) > 0:
|
|
value["description"] = description
|
|
target.append(value)
|
|
|
|
def _update_inputs(self, value, inputs):
|
|
target = value["inputs"] = []
|
|
for _ in inputs:
|
|
value = {}
|
|
value["name"] = _.name
|
|
value["type"] = _.typeStr
|
|
schemadef = onnxruntime.capi.onnxruntime_pybind11_state.schemadef
|
|
if _.option == schemadef.OpSchema.FormalParameterOption.Optional:
|
|
value["option"] = "optional"
|
|
elif _.option == schemadef.OpSchema.FormalParameterOption.Variadic:
|
|
value["list"] = True
|
|
description = _format_description(_.description)
|
|
if len(description) > 0:
|
|
value["description"] = description
|
|
target.append(value)
|
|
|
|
def _update_outputs(self, value, outputs):
|
|
target = value["outputs"] = []
|
|
for _ in outputs:
|
|
value = {}
|
|
value["name"] = _.name
|
|
value["type"] = _.typeStr
|
|
schemadef = onnxruntime.capi.onnxruntime_pybind11_state.schemadef
|
|
if _.option == schemadef.OpSchema.FormalParameterOption.Optional:
|
|
value["option"] = "optional"
|
|
elif _.option == schemadef.OpSchema.FormalParameterOption.Variadic:
|
|
value["list"] = True
|
|
description = _format_description(_.description)
|
|
if len(description) > 0:
|
|
value["description"] = description
|
|
target.append(value)
|
|
|
|
def _update_type_constraints(self, value, type_constraints):
|
|
value["type_constraints"] = []
|
|
for _ in type_constraints:
|
|
value["type_constraints"].append({
|
|
"description": _.description,
|
|
"type_param_str": _.type_param_str,
|
|
"allowed_type_strs": _.allowed_type_strs
|
|
})
|
|
|
|
def to_dict(self):
|
|
""" Serialize model to JSON message """
|
|
value = {}
|
|
value["name"] = self.name
|
|
value["module"] = self.module
|
|
value["version"] = self.version
|
|
schemadef = onnxruntime.capi.onnxruntime_pybind11_state.schemadef
|
|
if self.schema.support_level != schemadef.OpSchema.SupportType.COMMON:
|
|
value["status"] = self.schema.support_level.name.lower()
|
|
if self.schema.doc:
|
|
description = _format_description(self.schema.doc.lstrip())
|
|
if len(description) > 0:
|
|
value["description"] = description
|
|
if self.schema.attributes:
|
|
self._update_attributes(value, self.schema)
|
|
if self.schema.inputs:
|
|
self._update_inputs(value, self.schema.inputs)
|
|
value["min_input"] = self.schema.min_input
|
|
value["max_input"] = self.schema.max_input
|
|
if self.schema.outputs:
|
|
self._update_outputs(value, self.schema.outputs)
|
|
value["min_output"] = self.schema.min_output
|
|
value["max_output"] = self.schema.max_output
|
|
if self.schema.min_input != self.schema.max_input:
|
|
value["inputs_range"] = _format_range(self.schema.min_input) + " - " \
|
|
+ _format_range(self.schema.max_input)
|
|
if self.schema.min_output != self.schema.max_output:
|
|
value["outputs_range"] = _format_range(self.schema.min_output) + " - " \
|
|
+ _format_range(self.schema.max_output)
|
|
if self.schema.type_constraints:
|
|
self._update_type_constraints(value, self.schema.type_constraints)
|
|
return value
|
|
|
|
def _metadata():
|
|
root_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
|
|
file = os.path.join(root_dir, "source", "onnx-metadata.json")
|
|
with open(file, encoding="utf-8") as handle:
|
|
content = handle.read()
|
|
categories = {}
|
|
content = json.loads(content)
|
|
for schema in content:
|
|
if "category" in schema:
|
|
name = schema["name"]
|
|
categories[name] = schema["category"]
|
|
types = collections.OrderedDict()
|
|
numpy = __import__("numpy")
|
|
with numpy.errstate(all="ignore"):
|
|
snippets = onnx.backend.test.case.collect_snippets()
|
|
for schema in onnx.defs.get_all_schemas_with_history():
|
|
schema = OnnxSchema(schema, snippets)
|
|
if schema.key not in types:
|
|
types[schema.key] = schema.to_dict()
|
|
for schema in onnxruntime.capi.onnxruntime_pybind11_state.get_all_operator_schema():
|
|
schema = OnnxRuntimeSchema(schema)
|
|
if schema.key not in types:
|
|
types[schema.key] = schema.to_dict()
|
|
for schema in content:
|
|
key = f"{schema['name']}:{schema['module']}:{str(schema['version']).zfill(4)}"
|
|
if key not in types:
|
|
types[key] = schema
|
|
types = [types[key] for key in sorted(types)]
|
|
for schema in types:
|
|
name = schema["name"]
|
|
# copy = schema.copy()
|
|
# schema.clear()
|
|
# schema['name'] = name
|
|
# schema['module'] = copy['module']
|
|
if name in categories:
|
|
schema["category"] = categories[name]
|
|
# for key, value in copy.items():
|
|
# if key not in schema:
|
|
# schema[key] = value
|
|
content = json.dumps(types, indent=2)
|
|
with open(file, "w", encoding="utf-8") as handle:
|
|
handle.write(content)
|
|
|
|
def main():
|
|
_metadata()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|