自动微分在神经网络的向后反馈等机器学习算法时常用到。
梯度计算
当需要反馈时,TensorFlow需要追踪操作是以什么顺序在执行,以便于进行自动微分。反馈过程中,TensorFlow倒序遍历操作以计算梯度。
Gradient tapes
TensorFlow提供了tf.GradientTape的API进行自动微分,通常是对tf.Variables求梯度。TensorFlow会记录下tf.GradientTape上下文中的相关计算,然后通过倒序求解梯度。当使用tf.GradientTape记录下操作之后,就可以使用GradientTape.gradient(target, sources)来计算目因变量(通常是一个loss值)对于自变量的导数(通常是一个变量)
下面是一个小例子:
x = tf.Variable(3.0)
with tf.GradientTape() as tape:
y = x**2
dy_dx = tape.gradient(y, x)
dy_dx.numpy()
结果为:
6.0
尽管上面的小例子使用了变量是一个标量,但是tf.GradientTape可以作用于任意张量上。
w = tf.Variable(tf.random.normal((3, 2), name='w'))
b = tf.Variable(tf.zeros(2, dtype=tf.float32), name='b')
x = [[1., 2., 3.]]
with tf.GradientTape(persistent=True) as tape:
y = x @ w + b
loss = tf.reduce_mean(y ** 2)
上述代码中,为了同时求解loss对w,b,你可以将w和b同时传递给gradient函数。tape的使用十分灵活,你可以将多个变量组合为一个list或一个dict传递给gradient函数,那么gradient也会按照你的传递的格式进行返回。如下列代码所示:
[dl_dw, dl_db] = tape.gradient(loss, [w, b])
print(dl_dw.shape)
或
my_vars = {'w': w, 'b': b}
grad = tape.gradient(loss, my_vars)
print(grad['b'])
模型中的梯度计算
多数情况下,TensorFlow会收集tf.Module或其子类(tf.Model,tf.Layer)中的变量用于检查点设置和模型导出。你应该经常需要对模块(tf.Module)的变量(Moudle.trainable_variables)进行求导,这是几行代码就是能搞定的事情。
layer = tf.keras.layers.Dense(2, activation='relu')
x = tf.constant([[1., 2., 3.]])
with tf.GradientTape() as tape:
y = layer(x)
loss = tf.reduce_mean(y ** 2)
grad = tape.gradient(loss, layer.trainable_variables)
for var, g in zip(layer.trainable_variables, grad):
print(f'{var.name}, shape:{g.shape}')
上述代码的运算结果如下,代码没有什么实际意义,只是高速你变量和梯度的shape是一样的而已。
dense/kernel:0, shape:(3, 2)
dense/bias:0, shape:(2,)
控制tape监控的范围
tape默认情况下会记录tf.Variable的所有的相关计算操作。之所以这么设计,原因是:
- tape需要知道所有计算的顺序才能逆向求解微分。
- tape需要记录所有的中间运行结果,因此省却了使用者的麻烦。
- loss对所有的tf.Variable进行微分是用户最常用的操作。
下列代码中,微分只会对x1生效:
x0 = tf.Variable(3.0, name='x0')
x1 = tf.Variable(3.0, name='x1', trainable=False)
x2 = tf.Variable(2.0, name='x2') + 1.0
x3 = tf.constant(3.0, name='x3')
with tf.GradientTape() as tape:
y = x0 ** 2 + x1 ** 2 + x2 ** 2
grad = tape.gradient(y, [x0, x1, x2, x3])
for g in grad:
print(g)
运行结果为:
tf.Tensor(6.0, shape=(), dtype=float32)
None
None
None
注意到,x2实际上是一个张量,值为:
tf.Tensor(3.0, shape=(), dtype=float32)
tf.Gradient提供了让用户自由控制tape监控范围的功能。当需要对一个张量求微分时,你可以调用GradientTape.watch(x):
x = tf.constant(3.0)
with tf.GradientTape() as tape:
tape.watch(x)
y = x ** 2
dy_dx = tape.gradient(y, x)
print(dy_dx.numpy())
结果为
6.0
tape有可以修改默认的监控所有变量的行为。使用wach_accessed_variables=False可以关闭tape的默认行为,转而通过watch来自定义设置需要监控的变量。如下的代码中运算中使用了x0和x1两个变量,而微分过程只对x1进行。
x0 = tf.Variable(0.0)
x1 = tf.Variable(10.0)
with tf.GradientTape(watch_accessed_variables=False) as tape:
tape.watch(x1)
y0 = tf.math.sin(x0)
y1 = tf.nn.softplus(x1)
y = y0 + y1
ys = tf.reduce_sum(y)
grad = tape.gradient(ys, {'x0': x0, 'x1': x1})
print("dy_dx0:", grad['x0'])
print("dy_dx1:", grad['x1'])
运行结果为:
dy_dx0: None
dy_dx1: tf.Tensor(0.9999546, shape=(), dtype=float32)
中间结果
你可以使用tf.GradientTape来对中间的计算变量进行求微分。默认情况下,GradientTape所拥有的资源会在GradientTape.gradient之后释放内存,因此为了多次调用gradient方法以多次求微分,就需要使用persist=True参数来保存tape的监控内容。这种情况下,tape会随着Python生命作用域的消失而释放内存。
控制流
tape只记录执行过的操作,因此若是tape的代码里面有if-else,则tape只记录执行过的分支。
梯度计算出None的几种可能
- 无意中将变量替换成了一个张量
默认情况下,tape只会追踪监控tf.Variable,若是不经意间将tf.Variable变成了tf.Tensor,那么tape对其求微分便会是一个None值。因此应该使用Variable.assign方法来更新tf.Variable。
x = tf.Variable(2.0)
for epoch in range(2):
with tf.GradientTape() as tape:
y = x + 1
print(type(x).__name__, ":", tape.gradient(y, x))
x = x + 1
下面的例子中,x = x + 1使得tf.Variable变成了一个tf.Tensor。只需要保证在输入之前x是一个tf.Variable。
ResourceVariable : tf.Tensor(1.0, shape=(), dtype=float32)
EagerTensor : None
- 未使用TensorFlow提供的操作符进行计算
TensorFlow之外的计算并不能被tape所监控。例如:
x = tf.Variable([[1., 2.], [3., 4.]], dtype=tf.float32)
with tf.GradientTape() as tape:
x2 = x ** 2
y = np.mean(x2, axis=0)
y = tf.reduce_mean(y, axis=0)
print(tape.gradient(y, x))
结果为None,因为y=np.mean(x2, axis=0)这一行使用了numpy操作,不属于tensorFlow所提供的操作符。因此y与x之间没有了关联。
- 对整型或字符串求梯度
TensorFlow不能对整型和字符串求微分。开发时,用户自然不会考虑对字符串求微分,但是很有可能不经意间忘记指定dtype而创建了一个整型的张量或变量。
x = tf.constant(10)
with tf.GradientTape() as g:
g.watch(x)
y = x * x
print(g.gradient(y, x))
结果为空。
- 修改了变量的状态
状态的改变会导致梯度运算中断。当你需读取一个状态对象时,tape只能观察到当前的状态,而不会记录它的历史状态。因此,要注意在tape的记录过程中,不要修改变量的值。
下面这个程序的运行结果为None。
x0 = tf.Variable(3.0)
x1 = tf.Variable(4.0)
with tf.GradientTape() as tape:
x1.assign_add(x0)
y = x1**2
print(tape.gradient(y, x1))
再下面这个运行结果为14
x0 = tf.Variable(3.0)
x1 = tf.Variable(4.0)
with tf.GradientTape() as tape:
x1 = x1 + x0
y = x1**2
print(tape.gradient(y, x1))
下面程序的运行结果为:
x0 = tf.Variable(3.0)
x1 = tf.constant(4.0)
with tf.GradientTape() as tape:
x1 = x1 + 1
y = x1**2
print(tape.gradient(y, x1))
运行结果为:
None
有些操作是无法进行微分的
一些tf.Operation被注册为不可微分的,一旦计算了微分就会返回为None。还有一些连注册都没有,一旦试图对其求微分就会收到一个错误。tf.raw_ops页面记录了那些操作可以被微分。
将None替换为0
可以通过设置unconnnected_gradients将返回为None的操作的返回值重置为0。