将普通python函数转成tensorflow里的图（@tf.function）

tf.function：可以将普通的python函数转成tensorflow里的图
auto-graph：是tf.function所依赖的机制

将普通的python函数转成tensorflow里图结构的优势：

更快，tensorflow图结构里有很多的优化，所以一个普通的python函数转成tensorflow里的图实现的函数后，速度会提升。

除了使用tf.function来显式的把一个python函数转成tensorflow里的图，还可以用别的方法（添加@tf.function标注）把python函数转成tensorflow里的图，下面一一介绍。

1、 tf.function函数转换

核心代码展示：

# tf.function and auto-graph.
def scaled_elu(z, scale=1.0, alpha=1.0):
    # z >= 0 ? scale * z : scale * alpha * tf.nn.elu(z)
    is_positive = tf.greater_equal(z, 0.0) #判断z是否大于0
    return scale * tf.where(is_positive, z, alpha * tf.nn.elu(z)) #用where实现三项表达式

print(scaled_elu(tf.constant(-3.)))
print(scaled_elu(tf.constant([-3., -2.5])))

#一行代码将python语法实现的函数变成tensorflow中图实现的函数
scaled_elu_tf = tf.function(scaled_elu)

print(scaled_elu_tf(tf.constant(-3.)))
print(scaled_elu_tf(tf.constant([-3., -2.5])))

#对于转化后的函数可以通过  找回它的python函数
print(scaled_elu_tf.python_function is scaled_elu)

%timeit scaled_elu(tf.random.normal((1000, 1000)))
%timeit scaled_elu_tf(tf.random.normal((1000, 1000)))

运行结果为：

tf.Tensor(-0.95021296, shape=(), dtype=float32)
tf.Tensor([-0.95021296 -0.917915  ], shape=(2,), dtype=float32)
tf.Tensor(-0.95021296, shape=(), dtype=float32)
tf.Tensor([-0.95021296 -0.917915  ], shape=(2,), dtype=float32)
True
946 µs ± 12.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
838 µs ± 23.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

从上面的运行结果可以看出：1、普通python函数的运行结果与转换成tensorflow图实现的函数的运行结果相同

2、一个普通的python函数转成tensorflow里的图实现的函数后，速度会提升

2、@tf.function函数转换

核心代码展示：

# 1 + 1/2 + 1/2^2 + ... + 1/2^n

#@tf.function #标注：将下面的函数转成tensorflow里的图
def converge_to_2(n_iters): #迭代次数
    total = tf.constant(0.)
    increment = tf.constant(1.)
    for _ in range(n_iters):
        total += increment
        increment /= 2.0
    return total
print(converge_to_2(20))

未使用@tf.function输出结果为： 

tf.Tensor(1.9999981, shape=(), dtype=float32)

# 1 + 1/2 + 1/2^2 + ... + 1/2^n

@tf.function #标注：将下面的函数转成tensorflow里的图
def converge_to_2(n_iters): #迭代次数
    total = tf.constant(0.)
    increment = tf.constant(1.)
    for _ in range(n_iters):
        total += increment
        increment /= 2.0
    return total
print(converge_to_2(20))

使用@tf.function后输出结果为： 

tf.Tensor(1.9999981, shape=(), dtype=float32)

综上：用@tf.function也可以将一个普通的python函数转成tensorflow里的图实现的函数。

此外，使用to_code方法可以将python函数转成可以生成tensorflow图结构的中间代码展示出来，如下：

def display_tf_code(func): #传进来的是一个python函数
   """to_code函数可以将普通的Python函数转成tensorflow的代码，tensorflow的代码可以生成tensorflow里的图结构。"""
    code = tf.autograph.to_code(func)

    #展示tensorflow的代码
    from IPython.display import display, Markdown
    display(Markdown('```python\n{}\n```'.format(code)))

display_tf_code(scaled_elu)

输出结果为（输出的是转化成的tensorflow代码，scaled_elu函数在前面有定义）：

def tf__scaled_elu(z, scale=None, alpha=None):
  do_return = False
  retval_ = ag__.UndefinedReturnValue()
  with ag__.FunctionScope('scaled_elu', 'scaled_elu_scope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as scaled_elu_scope:
    is_positive = ag__.converted_call(tf.greater_equal, scaled_elu_scope.callopts, (z, 0.0), None, scaled_elu_scope)
    do_return = True
    retval_ = scaled_elu_scope.mark_return_value(scale * ag__.converted_call(tf.where, scaled_elu_scope.callopts, (is_positive, z, alpha * ag__.converted_call(tf.nn.elu, scaled_elu_scope.callopts, (z,), None, scaled_elu_scope)), None, scaled_elu_scope))
  do_return,
  return ag__.retval(retval_)

#使用变量,初始化必须在函数外面
var = tf.Variable(0.)

@tf.function
def add_21():
    return var.assign_add(21) # # 给变量+21+= 

print(add_21())

输出结果：

tf.Tensor(21.0, shape=(), dtype=float32)