Tensorflow-@tf_export

Apr 25, 2024
1 views
Python

@tf_export为函数取了个名字!


Tensorflow经常看到定义的函数前面加了@tf_export。例如,tensorflow/python/platform/app.py中有:

@tf_export('app.run')
def run(main=None, argv=None):
  """Runs the program with an optional 'main' function and 'argv' list."""

  # Define help flags.
  _define_help_flags()

  # Parse known flags.
  argv = flags.FLAGS(_sys.argv if argv is None else argv, known_only=True)

  main = main or _sys.modules['__main__'].main

  # Call the main function, passing through any arguments
  # to the final program.
  _sys.exit(main(argv))

首先,@tf_export是一个修饰符。修饰符的本质是一个函数

tf_export的实现在tensorflow/python/util/tf_export.py中:

```plain text tf_export = functools.partial(api_export, api_name=TENSORFLOW_API_NAME)


等号的右边的理解分两步:

1. functools.partial
1. api_export
functools.partial是偏函数,它的本质简而言之是为函数固定某些参数。如:`functools.partial(FuncA, p1)`的作用是把函数FuncA的第一个参数固定为p1;又如`functools.partial(FuncB, key1="Hello")`的作用是把FuncB中的参数key1固定为“Hello"。

`functools.partial(api_export, api_name=TENSORFLOW_API_NAME)`的意思是把api_export的api_name这个参数固定为TENSORFLOW_API。其中`TENSORFLOW_API_NAME = 'tensorflow'`。

api_export是实现了__call__()函数的类,简而言之是把类变得可以像函数一样调用。

`tf_export=functools.partial(api_export, api_name=TENSORFLOW_API_NAME)`的写法等效于:

```python
funcC = api_export(api_name=TENSORFLOW_API_NAME)
tf_export = funcC

对于funcC = api_export(api_name=TENSORFLOW_API_NAME),会导致api_export类的 __init__(api_name=TENSORFLOW_API_NAME)被调用:

  def __init__(self, *args, **kwargs):
    self._names = args
    self._names_v1 = kwargs.get('v1', args)
    self._api_name = kwargs.get('api_name', TENSORFLOW_API_NAME)
    self._overrides = kwargs.get('overrides', [])
    self._allow_multiple_exports = kwargs.get('allow_multiple_exports', False)

其中第4行self._api_name=kwargs.get('api_name', TENSORFLOW_API_NAME)的意思是获取api_name这个参数,如果未检测到该参数,则默认为TENSORFLOW_API_NAME。由此看,api_name这个参数传进来和默认的值都是TENSORFLOW_API_NAME,最终的结果是self._api_name=TENSORFLOW_API_NAME。

然后调用像函数一样调用funcC()实际上就会调用__call__():

def __call__(self, func):
    api_names_attr = API_ATTRS[self._api_name].names       #-----1
    api_names_attr_v1 = API_ATTRS_V1[self._api_name].names
    # Undecorate overridden names
    for f in self._overrides:
      _, undecorated_f = tf_decorator.unwrap(f)
      delattr(undecorated_f, api_names_attr)
      delattr(undecorated_f, api_names_attr_v1)

    _, undecorated_func = tf_decorator.unwrap(func)       #-----2
    self.set_attr(undecorated_func, api_names_attr, self._names)    #-----3
    self.set_attr(undecorated_func, api_names_attr_v1, self._names_v1)
    return func

因此@tf_export("app.run")最终的结果是用上面这个__call__()来作为修饰器。这是一个带参数的修饰器!

标注1:

api_names_attr = API_ATTRS[self._api_name].names中的self._api_name即为__init__()中提到的TENSORFLOW_API_NAME。看看API_ATTRS中都有些什么:

_Attributes = collections.namedtuple(
    'ExportedApiAttributes', ['names', 'constants'])

# Attribute values must be unique to each API.
API_ATTRS = {
    TENSORFLOW_API_NAME: _Attributes(
        '_tf_api_names',
        '_tf_api_constants'),
    ESTIMATOR_API_NAME: _Attributes(
        '_estimator_api_names',
        '_estimator_api_constants')
}

collections.namedtuple()返回具有命名字段的元组的新子类。从“ExportedApiAttributes”可推测这是用来管理已输出的API的属性的。

标注2:

_, undecorated_func = tf_decorator.unwrap(func)

def unwrap(maybe_tf_decorator):
  """Unwraps an object into a list of TFDecorators and a final target.
  Args:
    maybe_tf_decorator: Any callable object.
  Returns:
    A tuple whose first element is an list of TFDecorator-derived objects that
    were applied to the final callable target, and whose second element is the
    final undecorated callable target. If the `maybe_tf_decorator` parameter is
    not decorated by any TFDecorators, the first tuple element will be an empty
    list. The `TFDecorator` list is ordered from outermost to innermost
    decorators.
  """
  decorators = []
  cur = maybe_tf_decorator
  while True:
    if isinstance(cur, TFDecorator):
      decorators.append(cur)
    elif hasattr(cur, '_tf_decorator'):
      decorators.append(getattr(cur, '_tf_decorator'))
    else:
      break
    cur = decorators[-1].decorated_target
  return decorators, cur

将对象展开到tfdecorator列表和最终目标列表中。undecorated_func获得的返回对象就是我们有@tf_export修饰的函数。

标注3:self.set_attr(undecorated_func, api_names_attr, self._names) 设置属性。

总结:@tf_export修饰器为所修饰的函数取了个名字!

e.g. 通过@tf_export()就可以给函数重命名,比如@tf_export("nn.max_pool"),然后就可以使用tf.nn.max_pool()来调用tensorflow.python.ops.nn_ops.max_pool()函数