Hetu/python/hetu/onnx/X2hetu/handler.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import copy
import inspect
import os
import shutil

from onnx import defs


class BackendHandler:

    ONNX_OP = None

    DOMAIN = defs.ONNX_DOMAIN
    VERSION = 0
    SINCE_VERSION = 0
    PARTIAL_SUPPORT = False
    PS_DESCRIPTION = ""
    ONEFLOW_BLOBNAME_MAP = {}
    ONEFLOW_CODE_GEN = []
    OP_OUTPUS = []

    @classmethod
    def check_cls(cls):
        if not cls.ONNX_OP:
            print('doesn`t have ONNX_OP')

    @staticmethod
    def onnx_op(op):
        return BackendHandler.property_register("ONNX_OP", op)

    @classmethod
    def handle(cls, node, tensor_dict, **kwargs):

        ver_handle = getattr(cls, "version_{}".format(cls.SINCE_VERSION), None)
        if ver_handle:
            return ver_handle(node, tensor_dict, **kwargs)
        raise ValueError(
            'node "{}" of version {} is not supported'.format(
                node.op_type, cls.SINCE_VERSION
            )
        )

    @classmethod
    def run_onnx_node(cls,
                      node, tensor_dict,
                      inputs=None,
                      attrs=None,
                      name='',
                      **kwargs,
                      ):
        if inputs is None:
            inputs = [tensor_dict.get(inp, None)
                      for inp in node.input_tensor_names]
        if attrs is None:
            attrs = copy.deepcopy(node._attrs)
        if name != "":
            attrs["name"] = name
        for inp in node.input_tensor_names:
            if tensor_dict[inp] not in cls.ONEFLOW_BLOBNAME_MAP:
                cls.ONEFLOW_BLOBNAME_MAP[tensor_dict[inp]] = inp
        cls.OP_OUTPUS = []
        for oup in node.output_tensor_names:
            cls.OP_OUTPUS.append(oup)
        # todo
        # y = cls._run_flow_func(flow_func, inputs, attrs)
        # if type(y) == list():
        #     for x in cls.OP_OUTPUS:
        #         if y[x] not in cls.ONEFLOW_BLOBNAME_MAP:
        #             cls.ONEFLOW_BLOBNAME_MAP[y[x]] = x
        # else:
        #     if y not in cls.ONEFLOW_BLOBNAME_MAP:
        #         cls.ONEFLOW_BLOBNAME_MAP[y] = cls.OP_OUTPUS[0]
        return None  # y

    @staticmethod
    def property_register(name, value):
        def deco(cls):
            setattr(cls, name, value)
            return cls

        return deco


onnx_op = BackendHandler.onnx_op