wehub-resource-sync
7254f7b4d1
Build / Build (macos-latest) (push) Has been cancelled
Build / Build (ubuntu-latest) (push) Has been cancelled
Build / Build (windows-latest) (push) Has been cancelled
Build / Analyze (javascript) (push) Has been cancelled
Build / Analyze (python) (push) Has been cancelled
305 lines
11 KiB
Python
305 lines
11 KiB
Python
""" PyTorch backend """
|
|
|
|
import json
|
|
import os
|
|
|
|
|
|
class ModelFactory:
|
|
""" PyTorch backend model factory """
|
|
def open(self, model):
|
|
metadata = {}
|
|
metadata_files = [
|
|
("pytorch-metadata.json", ""),
|
|
("onnx-metadata.json", "onnx::")
|
|
]
|
|
path = os.path.dirname(__file__)
|
|
for entry in metadata_files:
|
|
file = os.path.join(path, entry[0])
|
|
with open(file, encoding="utf-8") as handle:
|
|
for item in json.load(handle):
|
|
name = entry[1] + item["name"].split("(", 1)[0]
|
|
metadata[name] = item
|
|
metadata = Metadata(metadata)
|
|
return _Model(metadata, model)
|
|
|
|
class _Model:
|
|
def __init__(self, metadata, model):
|
|
self.graph = _Graph(metadata, model)
|
|
|
|
def to_json(self):
|
|
""" Serialize model to JSON message """
|
|
import torch
|
|
json_model = {
|
|
"signature": "netron:pytorch",
|
|
"format": "TorchScript v" + torch.__version__,
|
|
"graphs": [ self.graph.to_json() ]
|
|
}
|
|
return json_model
|
|
|
|
class _Graph:
|
|
|
|
def __init__(self, metadata, model):
|
|
self.metadata = metadata
|
|
self.param = model
|
|
self.value = model.graph
|
|
self.nodes = []
|
|
|
|
def _getattr(self, node):
|
|
if node.kind() == "prim::Param":
|
|
return (self.param, "")
|
|
if node.kind() == "prim::GetAttr":
|
|
name = node.s("name")
|
|
obj, parent = self._getattr(node.input().node())
|
|
value = getattr(obj, name)
|
|
path = parent + "." + name if len(parent) > 0 else name
|
|
return (value, path)
|
|
raise NotImplementedError()
|
|
|
|
def to_json(self):
|
|
import torch
|
|
graph = self.value
|
|
json_graph = {
|
|
"values": [],
|
|
"nodes": [],
|
|
"inputs": [],
|
|
"outputs": []
|
|
}
|
|
data_type_map = dict([
|
|
[ torch.float16, "float16"],
|
|
[ torch.float32, "float32"],
|
|
[ torch.float64, "float64"],
|
|
[ torch.int32, "int32"],
|
|
[ torch.int64, "int64"],
|
|
])
|
|
def constant_value(node):
|
|
if node.hasAttribute("value"):
|
|
selector = node.kindOf("value")
|
|
return getattr(node, selector)("value")
|
|
return None
|
|
values_index = {}
|
|
def argument(value):
|
|
if value not in values_index:
|
|
json_value = {}
|
|
json_value["name"] = str(value.unique())
|
|
node = value.node()
|
|
if node.kind() == "prim::GetAttr":
|
|
tensor, name = self._getattr(node)
|
|
if tensor is not None and len(name) > 0 and \
|
|
isinstance(tensor, torch.Tensor):
|
|
json_tensor_shape = {
|
|
"dimensions": list(tensor.shape)
|
|
}
|
|
tensor_type = {
|
|
"dataType": data_type_map[tensor.dtype],
|
|
"shape": json_tensor_shape
|
|
}
|
|
json_value["name"] = name
|
|
json_value["type"] = tensor_type
|
|
json_value["initializer"] = { "type": tensor_type }
|
|
elif node.kind() == "prim::Constant":
|
|
tensor = constant_value(node)
|
|
if tensor and isinstance(tensor, torch.Tensor):
|
|
json_tensor_shape = {
|
|
"dimensions": list(tensor.shape)
|
|
}
|
|
tensor_type = {
|
|
"dataType": data_type_map[tensor.dtype],
|
|
"shape": json_tensor_shape
|
|
}
|
|
json_value["type"] = tensor_type
|
|
json_value["initializer"] = { "type": tensor_type }
|
|
elif value.isCompleteTensor():
|
|
json_tensor_shape = {
|
|
"dimensions": value.type().sizes()
|
|
}
|
|
json_value["type"] = {
|
|
"dataType": data_type_map[value.type().dtype()],
|
|
"shape": json_tensor_shape
|
|
}
|
|
values = json_graph["values"]
|
|
values_index[value] = len(values)
|
|
values.append(json_value)
|
|
return values_index[value]
|
|
|
|
for value in graph.inputs():
|
|
if len(value.uses()) != 0 and value.type().kind() != "ClassType":
|
|
json_graph["inputs"].append({
|
|
"name": value.debugName(),
|
|
"value": [ argument(value) ]
|
|
})
|
|
for value in graph.outputs():
|
|
json_graph["outputs"].append({
|
|
"name": value.debugName(),
|
|
"value": [ argument(value) ]
|
|
})
|
|
constants = {}
|
|
for node in graph.nodes():
|
|
if node.kind() == "prim::Constant":
|
|
constants[node] = 0
|
|
|
|
lists = {}
|
|
for node in graph.nodes():
|
|
if node.kind() == "prim::ListConstruct":
|
|
if all(_.node() in constants for _ in node.inputs()):
|
|
for _ in node.inputs():
|
|
constants[_.node()] += 1
|
|
lists[node] = 0
|
|
|
|
def create_node(node):
|
|
identifier = node.schema()
|
|
schema, category = self.metadata.type(identifier)
|
|
json_node = {
|
|
"type": {
|
|
"name": node.kind(),
|
|
"category": category
|
|
},
|
|
"inputs": [],
|
|
"outputs": [],
|
|
"attributes": []
|
|
}
|
|
json_graph["nodes"].append(json_node)
|
|
for name in node.attributeNames():
|
|
selector = node.kindOf(name)
|
|
value = getattr(node, selector)(name)
|
|
json_attribute = {
|
|
"name": name,
|
|
"value": value
|
|
}
|
|
if torch.is_tensor(value):
|
|
json_node["inputs"].append({
|
|
"name": name,
|
|
"value": []
|
|
})
|
|
else:
|
|
json_node["attributes"].append(json_attribute)
|
|
|
|
for i, value in enumerate(node.inputs()):
|
|
arg = None
|
|
if schema and i < len(schema.arguments):
|
|
arg = schema.arguments[i]
|
|
parameter_name = arg.name if arg else "input"
|
|
real_type = arg.real_type if arg else None
|
|
input_node = value.node()
|
|
if input_node in constants:
|
|
if (real_type and real_type.kind() == "TensorType") or \
|
|
value.type().kind() == "TensorType":
|
|
json_node["inputs"].append({
|
|
"name": parameter_name,
|
|
"value": [ argument(value) ]
|
|
})
|
|
else:
|
|
json_attribute = {
|
|
"name": parameter_name,
|
|
"value": constant_value(input_node)
|
|
}
|
|
if real_type:
|
|
json_attribute["type"] = self._argument_type(real_type)
|
|
json_node["attributes"].append(json_attribute)
|
|
constants[input_node] = constants[input_node] + 1
|
|
continue
|
|
if input_node in lists:
|
|
value = [ constant_value(_.node()) for _ in input_node.inputs() ]
|
|
json_attribute = {
|
|
"name": parameter_name,
|
|
"value": value
|
|
}
|
|
json_node["attributes"].append(json_attribute)
|
|
lists[input_node] += 1
|
|
continue
|
|
if input_node.kind() == "prim::TupleUnpack":
|
|
continue
|
|
if input_node.kind() == "prim::TupleConstruct":
|
|
continue
|
|
json_node["inputs"].append({
|
|
"name": parameter_name,
|
|
"value": [ argument(value) ]
|
|
})
|
|
|
|
for i, value in enumerate(node.outputs()):
|
|
ret = schema.returns[i] if schema and i < len(schema.returns) else None
|
|
name = ret.name if ret else "output"
|
|
json_node["outputs"].append({
|
|
"name": name,
|
|
"value": [ argument(value) ]
|
|
})
|
|
|
|
for node in graph.nodes():
|
|
if node in lists:
|
|
continue
|
|
if node in constants:
|
|
continue
|
|
if node.kind() == "prim::GetAttr":
|
|
continue
|
|
create_node(node)
|
|
|
|
for node in graph.nodes():
|
|
if node.kind() == "prim::Constant" and \
|
|
node in constants and constants[node] != len(node.output().uses()):
|
|
create_node(node)
|
|
if node.kind() == "prim::ListConstruct" and \
|
|
node in lists and lists[node] != len(node.output().uses()):
|
|
create_node(node)
|
|
|
|
return json_graph
|
|
|
|
def _argument_type(self, value):
|
|
if value.kind() == "TensorType":
|
|
return "Tensor"
|
|
if value.kind() == "OptionalType":
|
|
element_type = self._argument_type(value.getElementType())
|
|
return f"{element_type}?"
|
|
if value.kind() == "ListType":
|
|
element_type = self._argument_type(value.getElementType())
|
|
size = str(value.size) if hasattr(value, "size") else ""
|
|
return f"{element_type}[{size}]"
|
|
if value.kind() == "DictType":
|
|
key_type = self._argument_type(value.getKeyType())
|
|
value_type = self._argument_type(value.getValueType())
|
|
return f"Dict({key_type}, {value_type})"
|
|
if value.kind() == "TupleType":
|
|
elements = []
|
|
for element in value.elements():
|
|
elements.append(self._argument_type(element))
|
|
return f"({', '.join(elements)})"
|
|
if value.kind() == "IntType":
|
|
return "int64"
|
|
if value.kind() == "SymIntType":
|
|
return "SymInt"
|
|
if value.kind() == "FloatType":
|
|
return "float32"
|
|
if value.kind() == "BoolType":
|
|
return "boolean"
|
|
if value.kind() == "StringType":
|
|
return "string"
|
|
if value.kind() == "NumberType":
|
|
return "Scalar"
|
|
if value.kind() == "ScalarTypeType":
|
|
return "ScalarType"
|
|
if value.kind() == "LayoutType":
|
|
return "Layout"
|
|
if value.kind() == "MemoryFormatType":
|
|
return "MemoryFormat"
|
|
if value.kind() == "DeviceObjType":
|
|
return "Device"
|
|
if value.kind() == "GeneratorType":
|
|
return "Generator"
|
|
if value.kind() == "VarType":
|
|
return value.annotation_str
|
|
raise NotImplementedError()
|
|
|
|
class Metadata:
|
|
|
|
def __init__(self, metadata):
|
|
self.types = metadata
|
|
|
|
def type(self, identifier):
|
|
if identifier == "(no schema)":
|
|
return (None, "")
|
|
key = identifier.split("(", 1)[0]
|
|
value = self.types.get(key)
|
|
category = value["category"] if value and "category" in value else ""
|
|
name, overload_name = key.split(".", 1) if key.find(".") > 0 else (key, "")
|
|
import torch
|
|
schema = torch._C._get_schema(name, overload_name)
|
|
return (schema, category)
|