ascend-tools/pt2tf/onnx-tensorflow/onnx_tf/handlers/backend_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 tensorflow as tf

from onnx_tf.common import IS_PYTHON3
from onnx_tf.common import get_data_format
from onnx_tf.common import get_perm_from_formats
from onnx_tf.common import supports_device
from .handler import Handler


class BackendHandler(Handler):
  """ This class is base backend handler class.
  All backend operator handler class MUST inherit this class.
  In backend, operator handler class's name should be pascal case of file name
  which should be snake case.
  Use ONNX operator name as class name.
  """

  TF_FUNC = None

  @classmethod
  def get_attrs_processor_param(cls):
    """ Get param for attrs processor.

    :return: Dict.
    """
    return {}

  @classmethod
  def _process_attrs(cls, attrs):
    """ Private method for processing attrs.
    Param for this processor got from `get_attrs_processor_param`.
    Param is dict contains two key: `default` and `raname`.
    First add default value to attrs if key does not exist.
    Second rename key to new key.

    For example:
      attrs = {"keep_dims": True}
      param = {"default": {"axis": 1},
               "rename": {"keep_dims": "keepdims"}}

      processed_attrs = {"axis": "1", "keepdims": True}

    :param attrs: Process target attrs.
    :return: Processed attrs.
    """
    param = {"rename": {}, "default": {}}
    param.update(cls.get_attrs_processor_param())

    for k, v in param["default"].items():
      attrs.setdefault(k, v)

    for k, new_k in param["rename"].items():
      if k in attrs:
        attrs[new_k] = attrs.pop(k)

    return attrs

  @classmethod
  def make_tensor_from_onnx_node(cls,
                                 node,
                                 tf_func=None,
                                 inputs=None,
                                 attrs=None,
                                 name="",
                                 c_first_cuda_only=False,
                                 c_last_only=False,
                                 **kwargs):
    """ Helper method to make tensor.

    :param node: OnnxNode object.
    :param tf_func: Callable Tf function. Default is cls.TF_FUNC.
    :param inputs: Inputs tensor. Default is got from node.inputs.
    :param attrs: Attributes. Default is node.attrs.
    :param name: Node name.
    :param c_first_cuda_only: If channel first is only supported by cuda.
    If true and not cuda, do pre and post transpose.
    :param c_last_only: If only channel last is support,
    do pre and post transpose.
    :param kwargs: Other args.
    :return: Tensor.
    """
    tensor_dict = kwargs.get("tensor_dict", {})
    tf_func = tf_func or cls.TF_FUNC
    if tf_func is None:
      raise RuntimeError("No Tensorflow function is given.")
    if inputs is None:
      inputs = [tensor_dict.get(inp, None) for inp in node.inputs]
    if attrs is None:
      attrs = copy.deepcopy(node.attrs)
    name = name or node.name
    if name != "":
      attrs["name"] = name

    if c_first_cuda_only and c_last_only:
      raise ValueError(
          "c_first_cuda_only and c_last_only can not both be True.")

    if c_first_cuda_only:
      return cls.c_first_cuda_only(tf_func, inputs, attrs)
    elif c_last_only:
      return cls.c_last_only(tf_func, inputs, attrs)

    return cls._run_tf_func(tf_func, inputs, attrs)

  @classmethod
  def c_first_cuda_only(cls, tf_func, inputs, attrs):
    """ Handle operator that channel first is only supported by CUDA.
    When using CPU, two transposes should be added.

    :param tf_func: Callable Tf function.
    :param inputs: Inputs tensor.
    :param attrs: Attributes.
    :return: Tensor.
    """
    support_cuda = supports_device("CUDA")
    if not support_cuda:
      return cls._tuck_transpose(tf_func, inputs, attrs)
    return cls._run_tf_func(tf_func, inputs, attrs)

  @classmethod
  def c_last_only(cls, tf_func, inputs, attrs):
    """ Handle operator that channel last only is supported.
    Add two transposes anyway.

    :param tf_func: Callable Tf function.
    :param inputs: Inputs tensor.
    :param attrs: Attributes.
    :return: Tensor.
    """
    storage_format, compute_format = get_data_format(len(inputs[0].get_shape()))
    compute_format = compute_format.replace("C", "") + "C"
    return cls._tuck_transpose(tf_func, inputs, attrs,
                               (storage_format, compute_format))

  @classmethod
  def _tuck_transpose(cls, tf_func, inputs, attrs, data_format=None):
    x = inputs[0]
    x_rank = len(x.get_shape())
    if not data_format:
      data_format = get_data_format(x_rank)
    pre_perm = get_perm_from_formats(data_format[0], data_format[1])
    post_perm = get_perm_from_formats(data_format[1], data_format[0])
    attrs["data_format"] = data_format[1]
    if pre_perm != list(range(x_rank)):
      x_t = tf.transpose(x, perm=pre_perm)
      y = cls._run_tf_func(tf_func, [x_t] + inputs[1:], attrs)
      y_t = tf.transpose(y, perm=post_perm)
      return y_t
    return cls._run_tf_func(tf_func, inputs, attrs)

  @classmethod
  def _run_tf_func(cls, tf_func, inputs, attrs):
    """ Run Tensorflow function.
    Use only acceptable attributes of function from attrs.

    :param tf_func: Tensorflow function.
    :param inputs: Inputs.
    :param attrs: Attributes.
    :return: Tensor.
    """
    if IS_PYTHON3:
      params = list(inspect.signature(tf_func).parameters.keys())
    else:
      # use closure to get args for function using decorator
      if tf_func.__closure__ is not None:
        while "__wrapped__" in tf_func.func_dict:
          tf_func = tf_func.func_dict["__wrapped__"]
        params = inspect.getargspec(tf_func).args
      else:
        params = inspect.getargspec(tf_func).args

    attrs = cls._process_attrs(attrs)
    attrs = {p: v for p, v in attrs.items() if p in params}
    kwargs = dict(zip(params, inputs))
    ambiguous_arguments = any(kwargs.get(p) is not None and v is not None
                              for p, v in attrs.items())
    if ambiguous_arguments:
      raise TypeError('Ambiguous arguments for {}()'.format(tf_func.__name__))
    kwargs.update((p, v) for p, v in attrs.items() if v is not None)
    return tf_func(**kwargs)