（一）预备知识 -- 2 TensorFlow 求梯度/导数

2. TensorFlow 求梯度/导数

TensorFlow 为自动微分提供了 tf.GradientTape API ，即：
根据某个函数的输入变量来计算它的导数。

原理：
1）Tensorflow 会把 ‘tf.GradientTape’ 上下文中执行的所有操作都记录在一个磁带上 (“tape”)。
2）然后，基于这个磁带和每次操作产生的导数，用反向微分法（“reverse mode differentiation”）来计算这些被“记录在案”的函数的导数。

2.1 常量（constant）

x = tf.constant(3.0)
with tf.GradientTape() as tape:
    # track constants
    tape.watch(x)
    y = x*x
dy_dx = tape.gradient(y, x)
dy_dx.numpy()

6.0

tape.watch(variable) 作用：对指定变量/常量进行追踪。

GradientTape 默认不对常量进行追踪：

x = tf.constant(3.0)
with tf.GradientTape() as tape:
    y = x*x
dy_dx = tape.gradient(y, x)
print(dy_dx)

None

2.2 变量（variable）

2.2.1 trainable=True

GradientTape 支持对 trainable 变量进行自动追踪：

x = tf.Variable(3.0, trainable=True)
with tf.GradientTape() as tape:
    y = x*x
dy_dx = tape.gradient(y, x)
dy_dx.numpy()

6.0

2.2.2 watch_accessed_variables=False

GradientTape 支持对所有 trainable 变量中的一部分进行自动追踪。

x = tf.Variable(3.0, trainable=True)
m = tf.Variable(2.0, trainable=True)
with tf.GradientTape(watch_accessed_variables=False) as tape:
    tape.watch(x)
    y = x*x
    n = m*m
print(tape.gradient(y, x).numpy())

6.0

x = tf.Variable(3.0, trainable=True)
m = tf.Variable(2.0, trainable=True)
with tf.GradientTape(watch_accessed_variables=False) as tape:
    tape.watch(x)
    y = x*x
    n = m*m
print(tape.gradient(n, m))

None

如上所示，x 和 m 均为 trainable 变量，若仅对 x 进行追踪，
则能求得 y 对 x 的导数值为6.0，而 n 对 m 的导数返回None。

2.2.3 高阶导数

高阶导数的求取，可以通过 GradientTape 的嵌套来实现。

x = tf.Variable(3.0, trainable=True)
with tf.GradientTape() as tape_2:
    with tf.GradientTape() as tape_1:
        y = x**3
    g_1 = tape_1.gradient(y, x)
g_2 = tape_2.gradient(g_1, x)
g_2.numpy()

18.0

其中，调用 GradientTape.gradient() 方法时， 默认立即释放 GradientTape 占用的资源。因此，默认不支持连续两个 GradientTape.gradient() 的调用。

具体执行报错如下：

x = tf.Variable(3.0, trainable=True)
with tf.GradientTape() as tape:
    y = x**3
g_1 = tape.gradient(y, x)
g_2 = tape.gradient(g_1, x)
g_2

RuntimeError: GradientTape.gradient can only be called once on non-persistent tapes.

2.2.4 persistent=True

对比2.2.3，若需要支持连续两个 GradientTape.gradient() 的调用，可以通过 GradientTape 的参数配置来实现：

x = tf.Variable(3.0, trainable=True)
with tf.GradientTape(persistent=True) as tape:
    tape.watch(x)
    y = x*x
    z = y*y
print(tape.gradient(y, x).numpy())
print(tape.gradient(z, x).numpy())

6.0
108.0

2.2.5 stop_recording()

作用：暂时停止 tape 的记录。
用法：

x = tf.Variable(3.0, trainable=True)
m = tf.Variable(2.0, trainable=True)
with tf.GradientTape(persistent=True) as tape:
    y = x*x
    with tape.stop_recording():
        print(tape.gradient(y, x).numpy())
        
    n = m*m
    with tape.stop_recording():
        print(tape.gradient(n, m).numpy())

6.0
4.0

优点：增加代码可读性，提高计算速度。

2.2.6 python 控制流

TensorFlow 支持对 python 控制流求梯度（导数）。

def f(x, y):
    output = 1.0
    for i in range(y):
        if i > 1 and i < 5:
            output = tf.multiply(output, x)
    return output

def grad(x, y):
    with tf.GradientTape() as t:
        t.watch(x)
        out = f(x, y)
    return t.gradient(out, x)

x = tf.convert_to_tensor(2.0)
assert grad(x, 6).numpy() == 12.0

如上，当y=6.0时，out=x**3，对x求导并代入x=2.0，得12.0。

参考

tf.GradientTape Explained for Keras Users
TensorFlow 官方教程：自动微分和梯度带
《动手学深度学习》(TF2.0版)