@slim.add_arg_scope的使用方法与作用

今天在看代码是发现了好多函数前添加了装饰器@slim.add_arg_scope，一般装饰器都是为了使装饰的目标函数可以快速的调用，使整体代码更加简洁，而@slim.add_arg_scope的作用也是如此；首先先讲一下python中的装饰器；

1.Python装饰器(Decorator)

装饰器(Decorators)是 Python 的一个重要部分。简单地说：其就是修改其他函数的功能的函数。有助于代码更简短，同时可以在不修改源代码的同时添加功能函数，下面先建立一个装饰器，如下面代码所示： 

def decorator(func):
 
    def Function():
        print("before executing func()")
        func()
        print("after executing func()")
    return Function
 
def decoration():
    print("decoration")
 
decoration()
decoration = decorator(decoration)
decoration()

装饰器其实就是封装一个函数，并且用这样或者那样的方式来修改它的功能。代码中decoration = decorator(decoration)传入的是decoration，而不是decoration()，这是因为传入函数名不会执行函数，而是被传递，并且可以赋值给别的变量；若是传入了decoration()，则会直接执行函数，并将结果赋给变量；代码里并没有使用 @ 符号，因为那只是一个简短的方式来生成一个被装饰的函数。

def decorator(func):
 
    def Function():
        print("before executing func()")
        func()
        print("after executing func()")
    return Function
@decorator
def decoration():
    print("decoration")
 
decoration()

如上所示就是使用@符号简化的装饰器，其功能与上面代码一样；

但是有个问题，当执行

print(decoration.__name__)

发现其结果已经变了，不是decoration，而是Function,很明显其函数名已经被改写了，故需要使用方法来防止其被改写；Python提供给了一个简单的函数来解决这个问题，那就是functools.wraps。

from functools import wraps
def decorator(func):
    @wraps(func)
    def Function():
        print("before executing func()")
        func()
        print("after executing func()")
    return Function
@decorator
def decoration():
    print("decoration")
 
decoration()

现在执行

print(decoration.__name__)

发现其函数名没有被修改。同时@wraps也是个装饰器，它可以接收一个参数，就像普通的函数一样。故装饰器可以和其他函数一样接收参数，例如下面代码：

from functools import wraps

def log_with_param(text):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            print('call %s():' % func.__name__)
            print('args = {}'.format(*args))
            print('log_param = {}'.format(text))
            return func(*args, **kwargs)

        return wrapper

    return decorator
    
@log_with_param("param")
def test_with_param(p):
    print(test_with_param.__name__)

test_with_param("param")

其可以接收参数,与普通函数一样；同时类也可以用来构建装饰器。

from functools import wraps
 
class logit(object):
    def __init__(self, logfile='out.log'):
        self.logfile = logfile
 
    def __call__(self, func):
        @wraps(func)
        def wrapped_function(*args, **kwargs):
            log_string = func.__name__ + " was called"
            print(log_string)
            # 打开logfile并写入
            with open(self.logfile, 'a') as opened_file:
                # 现在将日志打到指定的文件
                opened_file.write(log_string + '\n')
            # 现在，发送一个通知
            self.notify()
            return func(*args, **kwargs)
        return wrapped_function
 
    def notify(self):
        # logit只打日志，不做别的
        pass

@logit()
def myfunc1():
    pass

上面代码是使用类来构建装饰器，从而调用函数。

2 .@slim.add_arg_scope

前面就是装饰器的基础使用和构造了，从而可知@slim.add_arg_scope也是用来装饰目标函数的，这里首先要了解其基本结构

这里的这篇博文讲的不错https://blog.csdn.net/weixin_35653315/article/details/78160886；

_ARGSTACK = [{}]
_DECORATED_OPS = {}

_ARGSTACK被当作一个栈来使用, 用于存储所有的scope. 一个scope就是一个dict. 它的key是str(fn), 其中fn代表目标函数；value是一个字典, 存储着用户调用slim.arg_scope时传入的参数. 之所以用栈, 是因为with是可以嵌套使用的. 栈顶存储的是最内层最近的scope, 即当前scope.
例如,

import tensorflow as tf
slim = tf.contrib.slim
@slim.add_arg_scope
def fn(a, b, c=3):
    d = c + b
    print("a={}, b={}".format(a, b))
    return d

with slim.arg_scope([fn], a = 1) as sc:
    fn(b = 2)
    print(sc)

在上面的这段代码中, 执行到fn(b=2)时, _ARGSTACK的值为:

{'<function fn at 0x000001F551522E18>': {'a': 1}}

_DECORATED_OPS是一个dict, 用于存储已经被@add_arg_scope修饰过的函数信息. 它的key也是str(fn), value的生成机制却很难get到它的真实意图, 给我的感觉一是有错, 二是无用. 所以, 把_DECORATED_OPS换成一个set应该也是可以的.

仍然是上面那段代码, _DECORATED_OPS的值为:

{'<function fn at 0x000001F551522E18>': ()}

add_arg_scope代码如下

def add_arg_scope(func):
  """Decorates a function with args so it can be used within an arg_scope.
  ...
  """
  @functools.wraps(func)
  def func_with_args(*args, **kwargs):
      ...
  _add_op(func) # 在这里将函数放入`_DECORATED_OPS`里.
  ...
  return func_with_args

_add_op的操作真的很简单, 看看_DECORATED_OPS里有没有目标函数, 没有就放进去:

def _add_op(op):
  key_op = _key_op(op)
  if key_op not in _DECORATED_OPS:
    _DECORATED_OPS[key_op] = _kwarg_names(op)
def _key_op(op):
  return getattr(op, '_key_op', str(op))

func_with_args内容暂时略过, 后面会分析.

with arg_scope([],...)

也是只保留了最关键的代码:

@contextlib.contextmanager
def arg_scope(list_ops_or_scope, **kwargs):
  """Stores the default arguments for the given set of list_ops.
    ...
  """
  ...
    try:
      current_scope = _current_arg_scope().copy()
      ...
        if key_op in current_scope:
          # 如果对一个函数多次应用arg_scope, 则更新已经存储的参数.
          current_kwargs = current_scope[key_op].copy()
          current_kwargs.update(kwargs)
          current_scope[key_op] = current_kwargs
        else:
          # 如果是第一次对这个函数调用arg_scope, 则加进去.
          current_scope[key_op] = kwargs.copy()
      # 将当前scope放入栈顶.
      _get_arg_stack().append(current_scope)
      yield current_scope
    finally:
      _get_arg_stack().pop()

def _get_arg_stack():
  if _ARGSTACK:
    return _ARGSTACK
  else:
    _ARGSTACK.append({})
    return _ARGSTACK

在with代码块内部调用目标函数，之前省略的函数func_with_args其实是关键中的关键: 在with代码块内调用被@add_arg_scope修饰的函数, 其实是调用了它, 而它负责从当前栈中查找用户设置的默认参数, 加上用户传入的调用参数, 调用真正的源函数.
func_with_args的代码:

  def func_with_args(*args, **kwargs):
    current_scope = _current_arg_scope()# 当前栈
    current_args = kwargs # 用户传入的关键字参数
    key_func = _key_op(func)  
    if key_func in current_scope:
      # 通过arg_scope设置的默认参数
      current_args = current_scope[key_func].copy() 
      current_args.update(kwargs)
    # 调用原函数
    return func(*args, **current_args)

 

退出with模块

弹出栈顶的scope. 代码在arg_scope方法最后一行.