netron/source/pytorch.py

""" 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)