【TF2】Eager Execution机制

1. Eager Execution 简介

1.1 eager execution 引入

eager 模式是在 TF 1.4 版本之后引入的，并在TF2.0+将eager模式便成为默认执行模式

1.2 理解eager execution

自己理解，eager 模式就是类似于 python 这样的命令式编程，写好程序之后，不需要编译，就可以直接运行了(其实是一边建图一边计算，不具有复用性)，而且非常直观；而之前的静态图模式则类似于 c/c++ 的声明式编程。写好程序之后要先编译(搭建好静态计算图)，然后才能运行。

1.3 eager execution优缺点

优点

• eager 模式提供了更直观的接口，可以像写 python 代码一样写模型；
• 更方便调试，这个应该是不言而喻的，如果同时调试过 python 和 c++ 就知道有多方便了；
• 自然的控制流程，像编写 python 程序一样；

当然就像 python 一直被诟病的速度慢的问题以及不适合做大的工程的问题，eager 模式也具有类似的问题：

缺点

• 速度问题：对于有很大计算量的模型，比如在 GPU 上训练 ResNet50，eager 模式和通过 graph 模式构造的模型没有太大的差异，但是这个差异却随着计算量减小而扩大（也就是说，如果模型越简单，通过 eager 模式构造的模型速度上要比通过 graph 构造的模型慢）
• 分布式训练、性能优化以及线上部署：通过 graph 构造的模型在分布式训练、性能优化以及线上部署上有优势；

建议

eager 模式支持大多数的 TF 操作和 GPU 加速；TF 官方的一些例子可以参考：https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples

基于此，官方给的建议是：同时写 eager 模式代码和 graph 模式代码，使用 eager 模式快速调试和 debug，然后利用 graph 模式在分布式训练方面的优势部署；

2. eager execution用法

如果 TF 版本在 1.4 以下，使用如下方式进行升级：

pip install --upgrade tensorflow

2.1 开启eager模式

import tensorflow as tf
tf.enable_eager_execution()

tf.executing_eagerly()        # => True

x = [[2.]]
m = tf.matmul(x, x)
print("hello, {}".format(m))  # => "hello, [[4.]]"

eager 模式下，操作会立刻执行并返回实际值。在这里 tf.Tensor 对象会是一个实际值而不是 graph 模式下的 Tensor 结构（这个 Tensor 结构主要包含了三个属性：名字、维度和类型，而且只有运行 graph 的时候，Tensor 才会有真实的值）；

2.1 eager execution 和 numpy

eager 模式下对 numpy 的支持很友好。具体有如下三个特性：

• numpy 的操作可以接受 Tensor 作为参数；
• TF 的数学操作会将 python 对象和 numpy 的 arrays 转换成 Tensor；
• tf.Tensor.numpy 方法返回 numpy 的 ndarray；

直接看官方的例子即可：

a = tf.constant([[1, 2],
                 [3, 4]])
print(a)
# => tf.Tensor([[1 2]
#               [3 4]], shape=(2, 2), dtype=int32)

# Broadcasting support
b = tf.add(a, 1)
print(b)
# => tf.Tensor([[2 3]
#               [4 5]], shape=(2, 2), dtype=int32)

# Operator overloading is supported
print(a * b)
# => tf.Tensor([[ 2  6]
#               [12 20]], shape=(2, 2), dtype=int32)

# Use NumPy values
import numpy as np

c = np.multiply(a, b)
print(c)
# => [[ 2  6]
#     [12 20]]

# Obtain numpy value from a tensor:
print(a.numpy())
# => [[1 2]
#     [3 4]]

3. 动态控制流程

这是 eager 模式的主要特性，就是像写 python 代码一样写 TF 代码，而且可以随时通过 print 来进行调试；

一个 fizzbuzz 示例：

def fizzbuzz(max_num):
  counter = tf.constant(0)
  max_num = tf.convert_to_tensor(max_num)
  for num in range(max_num.numpy()):
    num = tf.constant(num)
    if int(num % 3) == 0 and int(num % 5) == 0:
      print('FizzBuzz')
    elif int(num % 3) == 0:
      print('Fizz')
    elif int(num % 5) == 0:
      print('Buzz')
    else:
      print(num)
    counter += 1
  return counter

4. 建模

模型一般由多层组成，你可以自定义层，也可以使用 tf.keras.layers 提供的 layers。或者继承 tf.keras.layers.Layer 基类来实现自己的 layer；

代码如下：

lass MySimpleLayer(tf.keras.layers.Layer):
  def __init__(self, output_units):
    super(MySimpleLayer, self).__init__()
    self.output_units = output_units

  def build(self, input_shape):
    # 在第一次使用 layer 的时候才会调用 build 方法，因此定义在 build 下面的 input_shape 可以动态指定，当然，如果你知道具体的 layer 大小，也可以在 __init__ 方法中直接指定；
    self.kernel = self.add_variable(
      "kernel", [input_shape[-1], self.output_units])

  def call(self, input):
    # Override call() instead of __call__ so we can perform some bookkeeping.
    return tf.matmul(input, self.kernel)

当然也可以使用 tf.keras.layers.Dense 和 tf.keras.Sequential 来组织模型：

model = tf.keras.Sequential([
tf.keras.layers.Dense(10, input_shape=(784,)), # must declare input shape
tf.keras.layers.Dense(10)
])

而且，你也可以使用 python 的类（继承 tf.keras.Model）配合 tf.keras.layers 建模；

class MNISTModel(tf.keras.Model):
  def __init__(self):
    super(MNISTModel, self).__init__()
    self.dense1 = tf.keras.layers.Dense(units=10)
    self.dense2 = tf.keras.layers.Dense(units=10)

  def call(self, input):
    """Run the model."""
    result = self.dense1(input)
    result = self.dense2(result)
    result = self.dense2(result)  # reuse variables from dense2 layer
    return result

model = MNISTModel()

5. Eager训练

5.1 梯度计算

在 eager 模式中，可以使用 tf.GradientTape 来跟踪记录操作；Tape 是磁带的意思，因此可以这样理解，在前向计算的时候，相关操作会记录到 Tape 上，然后反向播放 Tape 即可计算梯度；一个特定的 tf.GradientTape 只能计算一个梯度，再一次调用就会报错；

简单的示例：

w = tf.Variable([[1.0]])
# 前向计算，得到 loss
with tf.GradientTape() as tape:
  loss = w * w

grad = tape.gradient(loss, w)
print(grad)  # => tf.Tensor([[ 2.]], shape=(1, 1), dtype=float32)

5.2 变量和优化

可以用类将模型封装，模型的参数作为类的参数

以3*x+2

class Model(tf.keras.Model):
  def __init__(self):
    super(Model, self).__init__()
    self.W = tf.Variable(3., name='weight')
    self.B = tf.Variable(2., name='bias')
  def call(self, inputs):
    return inputs * self.W + self.B

6. 使用 python 对象来存储程序状态

在 graph 模式中，程序状态（比如变量）是被存储在一个全局集合中，生命周期是由 tf.Session 管理的。eager 模式中，程序状态的生命周期是由其对应的 python 对象来决定的

6.1 变量都是对象

eager 模式下，直到引用该变量的最后一个对象被移除，变量才会被删除；

示例代码：

with tf.device("gpu:0"):
  v = tf.Variable(tf.random_normal([1000, 1000]))
  v = None  # v no longer takes up GPU memory

6.2 基于对象的保存

tf.train.Checkpoint 可以保存和恢复 tf.Variable，tf.keras.Model等；
参考文章：https://blog.csdn.net/weixin_44441131/article/details/121626669

6.3 Summaries 和 TensorBoard

tf.contrib.summary 兼容 eager execution 模式和 graph 模式

比如，每 100 步记录一下 summaries

global_step = tf.train.get_or_create_global_step()
writer = tf.contrib.summary.create_file_writer(logdir)
writer.set_as_default()

for _ in range(iterations):
  global_step.assign_add(1)
  # Must include a record_summaries method
  with tf.contrib.summary.record_summaries_every_n_global_steps(100):
    # your model code goes here
    tf.contrib.summary.scalar('loss', loss)
     ...

7. 自动微分

7.1 动态模型

tf.GradientTape 可以用在动态模型中，下面的示例是：backtracking line search 算法，看起来就像普通的 numpy 代码，但是这里包含了梯度计算和自动微分：

def line_search_step(fn, init_x, rate=1.0):
  with tf.GradientTape() as tape:
    # Variables are automatically recorded, but manually watch a tensor
    tape.watch(init_x)
    value = fn(init_x)
  grad = tape.gradient(value, init_x)
  grad_norm = tf.reduce_sum(grad * grad)
  init_value = value
  while value > init_value - rate * grad_norm:
    x = init_x - rate * grad
    value = fn(x)
    rate /= 2.0
  return x, value

7.2 计算梯度的方法

除了 tf.GradientTape 方法之外，还有一些其他的自动微分的方法。如果直接计算数学公式的梯度（数学公式中只有 Tensor，没有 tf.Variables），那么这些方法会很有用；

7.2.1 tfe.gradients_function

• 输入是一个函数；
• 输出该函数的梯度函数（其实就是一个求导后的函数），给定输入这个梯度函数会自动计算梯度，返回的是一个 list，其中的每一个元素是输入参数对应的梯度；

注：这里说的有点绕，还是直接看代码吧：

def square(x):
  return tf.multiply(x, x)

grad = tfe.gradients_function(square)

square(3.)  # => 9.0
grad(3.)    # => [6.0]

def square2(x, y):
  return tf.multiply(x, x) + tf.multiply(y, y)

grad = tfe.gradients_function(square2)
square2(3., 2.) # => 13
grad(3., 2.)    # => 6, 4

# The second-order derivative of square:
gradgrad = tfe.gradients_function(lambda x: grad(x)[0])
gradgrad(3.)  # => [2.0]

# The third-order derivative is None:
gradgradgrad = tfe.gradients_function(lambda x: gradgrad(x)[0])
gradgradgrad(3.)  # => [None]


# With flow control:
def abs(x):
  return x if x > 0. else -x

grad = tfe.gradients_function(abs)

grad(3.)   # => [1.0]
grad(-3.)  # => [-1.0]

7.2.2 tfe.value_and_gradients_function

梯度函数输出的时候会带上 value（输入函数计算出来的），示例代码如下：

def square(x, y):
    return tf.multiply(x, x) + tf.multiply(y, y)

grad = tfe.value_and_gradients_function(square)
print(square(3., 2.))
print(grad(3., 2.))

示例输出：

tf.Tensor(13.0, shape=(), dtype=float32)
(<tf.Tensor: id=13, shape=(), dtype=float32, numpy=13.0>, [<tf.Tensor: id=19, shape=(), dtype=float32, numpy=6.0>, <tf.Tensor: id=20, shape=(), dtype=float32, numpy=4.0>])

7.3 定制梯度

@tf.custom_gradient 方式可以对梯度进行一下定制，通常的用法是提供一个更平滑的梯度（防止梯度爆炸），比如约束梯度的 l2 范数；
代码示例：

@tf.custom_gradient
def clip_gradient_by_norm(x, norm):
  y = tf.identity(x)
  def grad_fn(dresult):
    return [tf.clip_by_norm(dresult, norm), None]
  return y, grad_fn

上面这种方式会因为数值计算的不稳定性导致结果为 nan，通过定制梯度可以解决这种问题

@tf.custom_gradient
def log1pexp(x):
  e = tf.exp(x)
  def grad(dy):
    # 梯度规约
    return dy * (1 - 1 / (1 + e))
  return tf.log(1 + e), grad

grad_log1pexp = tfe.gradients_function(log1pexp)

# As before, the gradient computation works fine at x = 0.
grad_log1pexp(0.)  # => [0.5]

# And the gradient computation also works at x = 100.
grad_log1pexp(100.)  # => [1.0]

更多性能、和graph模式一起工作参考https://zhuanlan.zhihu.com/p/47201474