【Eager & Graph Execution】EagerTensor与Tensor，tf.function / tf.py_function / tf.numpy_function

1. Eager execution 与 Graph execution

Tensorflow 的执行方式可分为：

• Eager execution
• Graph execution, 也叫静态图执行方式

Graph execution早于Eager execution：

在TensorFlow v1.5之前，tensorflow都是通过计算图将计算的定义和执行分隔开, 这是一种声明式（declaretive）的编程模型。这种静态图的执行模式最主要的特点是快速。但是在debug时非常不方便，无法类似于对编译好的Python语言程序调用，也无法对其进行内部的调试。

因此有了Eager Execution，这在TensorFlow v1.5首次引入。(最开始我以为是tensorflow2.x才有的执行方式)。引入Eager Execution模式后, TensorFlow就拥有了类似于Pytorch一样动态图模型能力, 我们可以不必再等到sess.run(*)才能看到执行结果, 可以方便在IDE随时调试代码,查看每一行的执行结果。

两种方式各有优点，用户的需求可能是既能静态图运行某些部分，又能Eager模式进行Debug：

那么可以在静态图执行方式中，启动 Eager Execution：（具体是不是tensorflow1.5版本之前都可以这样启动Eager Execution待确定）

import tensorflow as tf # 2.4.1
tf.compat.v1.enable_eager_execution()

也希望能在Eager Execution中，进行静态图执行：

@tf.function
def our_functions():
	tf.print("This is graph execution!")

ref：

tensorflow图模式和eager模式小结: https://zhuanlan.zhihu.com/p/82548504:
tf2通过tf.function这个API来做到最小化编程差异，你可以首先写好eager模式下的代码，debug完之后，使用tf.function修饰你想用图的方式运行的函数，那么这个函数就会使用图的模式运行。

2. EagerTensor 与 Tensor

Tensor：tensorflow.python.framework.ops.Tensor
EagerTensor： tensorflow.python.framework.ops.EagerTensor

EagerTensor属于Tensor

很多时候也会把EagerTensor直接显示为tensor，如：

 a = tf.constant([1,2,3,4,5])
 
a
<tf.Tensor: shape=(5,), dtype=int32, numpy=array([1, 2, 3, 4, 5])>
 
type(a)
tensorflow.python.framework.ops.EagerTensor

TF2.x默认Tensor都是EagerTensor

或者在eager execution下的都是eager tensor:

import tensorflow as tf
print(tf.__version__) # 2.4.1

def f(x, y):
  return x ** 2 + y

if __name__ == '__main__':
    out = f(tf.constant(2), tf.constant(3))

虽然打印出来依然是tf.Tensor:

tf.Tensor(7, shape=(), dtype=int32)

EagerTensor 和 Tensor 区别

获取数值: EagerTensor.numpy(), Tensor.eval()

• EagerTensor有numpy()属性， 而Tensor则没有。所以静态图运行下的Tensor无法通过numpy()获取数值。

  tensorflow1.x 静态图中，我们需要采用以下方式来获取Tensor的值：

  with tf.Session() as sess:
      result = sess.run([out])  # 获取Tensor out 的值
      print(result)
      
      # 或者是
      result = out.eval()
  

打印：tf.print(EagerTensor/Tensor), print(EagerTensor)

• tf.print(): tf.print(tensor)能够无论在动态图还是静态图下都能够打印出张量的内容，即既可以tf.print(tensor)，也可以tf.print(EagerTensor)。

• print(): 而print(EagerTensor.numpy())只能在动态图下即Eager Execution使用，而且只能够对EagerTensor使用。

• EagerTensor可以直接与python代码进行交互的，也可以进行迭代遍历操作，而Tensor则不支持与Python直接进行交互。

3. @tf.function：Eager execution 转 Graph execution

即Eager execution转化为静态图执行方式Graph execution。

Ref:
https://tensorflow.google.cn/api_docs/python/tf/function

声明与描述

声明：

tf.function(
    func=None, input_signature=None, autograph=True, experimental_implements=None,
    experimental_autograph_options=None, experimental_relax_shapes=False,
    experimental_compile=None, experimental_follow_type_hints=None
)

描述：
tf.function constructs a callable that executes a TensorFlow graph (tf.Graph) created by trace-compiling the TensorFlow operations in func, effectively executing func as a TensorFlow graph.
即以静态图形式调用

调用方式: 函数前加@tf.function

在我们写好并debug好我们的eager function后，可通过在函数前加@tf.function，将该函数以静态图方式运行：非常方便地将函数静态化。此时该函数不再可被debug，即无法被breakpoint捕获。

官方示例：

@tf.function
def f(x, y):
  return x ** 2 + y
  
x = tf.constant([2, 3])
y = tf.constant([3, -2])

z = f(x, y)

将某一函数以静态图形式调用：

• @tf.function

定义tensor变量及调用函数：

• f(x, y)

输入输出

@tf.function修饰的静态函数输入可为:

• int，numpy（即ndarray），tf.Variable（即ResourceVariable）, 或tf.Tensor（即tf.EagerTensor）

但输出都为

• tf.Tensor（即tf.EagerTensor）:

import numpy as np
import tensorflow as tf
print(tf.__version__)

@tf.function
def f(x, y):
  return x ** 2 + y

if __name__ == '__main__':
    x = 2
    y = 3
    out = f(x, y)
    print(out)

    x = np.ones(2)
    y = np.ones(2)
    out1 = f(x, y)
    print(out1)

    x = tf.constant(2)
    y = tf.constant(3)
    out2 = f(x, y)
    print(out2)

	x = tf.Variable(2)
    y = tf.Variable(3)
    out3 = f(x, y)
    print(out3)

运行结果：

• 输入：
  
  
  
  

• 输出：
  

特点

• func may use data-dependent control flow, including if, for, while break, continue and return statements。即虽然经过@tf.function修饰的函数，将在静态图下运行，但是依然可以用if，for，while break等，而不用tf.cond，tf.while_loop等。

• func may also use ops with side effects, such as tf.print, tf.Variable and others, e.g.,

  import tensorflow as tf
  print(tf.__version__) # 2.4.1
  
  v = tf.Variable(1) # 不能放在函数f_side_effect里
  
  @tf.function
  def f_side_effect(x):
      for i in tf.range(x):
          v.assign_add(i)
          
          # 请注意这两个print
          print('print by print') 
          tf.print('print by tf.print')
      return v
      
  if __name__ == '__main__':
  	out = f_side_effect(3)
  

  输入：
  
  输出：
  

• Key Point: Any Python side-effects (appending to a list, printing with print, etc) will only happen once, when func is traced. To have side-effects executed into your tf.function they need to be written as TF ops:

  即类似于list appending，printing这样的side-effects的操作，只会仅仅执行一次。需以TF的方式来写，如print()需写为tf.print()。其实从print()和tf.print()本身的打印特点也可知，print(EagerTensor)，tf.print(EagerTensor/Tensor)。

  • e.g., print() 与 tf.print()

  import tensorflow as tf
  print(tf.__version__) # 2.4.1
  
  v = tf.Variable(1) # 不能放在函数f_side_effect里
  
  @tf.function
  def f_side_effect(x):
      for i in tf.range(x):
          v.assign_add(i)
          
          # 请注意这两个print
          print('print by print') 
          tf.print('print by tf.print')
      return v
      
  if __name__ == '__main__':
  	out = f_side_effect(3)
  

  打印结果：

  

• Internally, tf.function can build more than one graph, to support arguments with different data types or shapes, since TensorFlow can build more efficient graphs that are specialized on shapes and dtypes. tf.function also treats any pure Python value as opaque objects, and builds a separate graph for each set of Python arguments that it encounters.

  即tf.function在输入类型或shape不同时，会创建不同的tf.Graph。并且会为每个python 对象分别创建静态图，如：

  @tf.function
  def f(x):
  	return x + 1
  	
  # 可通过 f.get_concrete_function(1).graph获取静态图  
  >>> isinstance(f.get_concrete_function(1).graph, tf.Graph)
  >>> True
  
  
  
  ##########################################################
  @tf.function
  def f(x):
  	return tf.abs(x)
  
  # 输入int
  f1 = f.get_concrete_function(1)
  f2 = f.get_concrete_function(2)  # Slow - builds new graph
  >>> f1 is f2
  >>> False
  
  # 输入Tensor
  f1 = f.get_concrete_function(tf.constant(1))
  f2 = f.get_concrete_function(tf.constant(2))  # Fast - reuses f1
  >>> f1 is f2
  >>> True
  
  
  ##########################################################
  @tf.function
  def f(x):
  	return x + 1
  
  # 输入shape不同的Tensor
  vector = tf.constant([1.0, 1.0])
  matrix = tf.constant([[3.0]])
  >>> f.get_concrete_function(vector) is f.get_concrete_function(matrix)
  >>> False
  

4. tf.py_function：Graph execution 转 Eager execution

反之，也有在Graph execution中执行Eager execution：tf.py_function

Ref:
tf1.15: https://tensorflow.google.cn/versions/r1.15/api_docs/python/tf/py_function,
tf2.4.1: https://tensorflow.google.cn/api_docs/python/tf/py_function

声明与描述

声明：

tf.py_function(
    func, inp, Tout, name=None
)

描述：

Wraps a python function into a TensorFlow op that executes it eagerly:
即将普通的 python函数包装成Tensorflow的操作，以Eager Execution的执行方式执行。

This function allows expressing computations in a TensorFlow graph as Python functions. In particular, it wraps a Python function func in a once-differentiable TensorFlow operation that executes it with eager execution enabled. As a consequence, tf.py_function makes it possible to express control flow using Python constructs (if, while, for, etc.), instead of TensorFlow control flow constructs (tf.cond, tf.while_loop).
同时tf.py_function 有一个好处是可以使用python的结构，如if、while break、for等，而不需要tensorflow的结构如tf.cond, tf.while_loop。

调用方式：tf.py_function(func=函数名，inp=输入参数，Tout=输出参数类型)

使用 tf.py_function 需要指定形状和类型信息，否则它将不可用。

官方示例：

For example, you might use tf.py_function to implement the log huber function:

def log_huber(x, m):
  if tf.abs(x) <= m:
    return x**2
  else:
    return m**2 * (1 - 2 * tf.math.log(m) + tf.math.log(x**2))

x = tf.compat.v1.placeholder(tf.float32)
m = tf.compat.v1.placeholder(tf.float32)

y = tf.py_function(func=log_huber, inp=[x, m], Tout=tf.float32)
dy_dx = tf.gradients(y, x)[0]

with tf.compat.v1.Session() as sess:
  # The session executes `log_huber` eagerly. Given the feed values below,
  # it will take the first branch, so `y` evaluates to 1.0 and
  # `dy_dx` evaluates to 2.0.
  y, dy_dx = sess.run([y, dy_dx], feed_dict={x: 1.0, m: 2.0})

定义一个huber函数：

• def log_huber(x, m)

以tf.compat.v1.placeholder 初始化tensorflow 1.x (<1.5)的tensor（此此处所用的tf应该是>=1.5）:

• x = tf.compat.v1.placeholder(tf.float32)

再以tf.py_function 调用huber函数:

• y = tf.py_function(func=log_huber, inp=[x, m], Tout=tf.float32)

启动sess进行静态图计算：

• sess.run(y, feed_dict)

与tf.py_func区别

• tf.py_func: Wraps a python function and uses it as a TensorFlow op. (https://tensorflow.google.cn/versions/r1.15/api_docs/python/tf/py_func), tf.py_func 只有tf1.x版本。

• tf.py_function: Wraps a python function into a TensorFlow op that executes it eagerly.

• tf.py_function is similar in spirit to tf.compat.v1.py_func, but unlike the latter, the former lets you use TensorFlow operations in the wrapped Python function. In particular, while tf.compat.v1.py_func only runs on CPUs and wraps functions that take NumPy arrays as inputs and return NumPy arrays as outputs, tf.py_function can be placed on GPUs and wraps functions that take Tensors as inputs, execute TensorFlow operations in their bodies, and return Tensors as outputs.

• 即tf.py_func只能跑在CPUs，且只能输入输出Numpy

• tf.py_function可以跑在GPUs，可以以Tensor为输入输出

应用示例

在tf.data.Dataset.map时，使用 tf.data.Dataset.map 方法可将函数应用于 Dataset 的每个元素。映射函数必须在 TensorFlow 计算图模式下进行运算（它必须在 tf.Tensors 上运算并返回）。

如该map函数是一个非张量函数，或包含一些非张量变量：这时，可以使用 tf.py_function 包装非张量函数（如 serialize_example）以使其兼容。

总结：

map_function默认是使用静态graph来执行，所以传入给map_function的参数都会自动被包装成Tensor，要想在map_function里面编写解析的Python代码：

通过tf.py_function()来包装map_function函数。

注意一点，一个自定义的解析函数，通过dataset.map(func)，之后，不管解析的数据本身是EagerTensor还是Tensor，都会变成Tensor，因为map函数会将它所包装的func转化成静态graph，所以总是Tensor

5. tf.numpy_function

它所包装的函数只能够接受numpy.array作为函数参数，并且返回的也是numpy.array，这里就不再赘述了。