深度学习笔记-tf.GradientTape用法详解

TensorFlow GradientTape 用法详解

最近看到莫凡老师上传的GAN相关的代码时, 发现莫凡老师写的GAN网络的代码使用了tf.GradientTape()进行先计算模型的梯度,然后利用优化器的apply_gradients()方法进行训练优化,于是查了一下相应的api总结出了此篇文档.

1. 简单的梯度计算

import tensorflow as tf
x = tf.Variable(3.)
with tf.GradientTape() as tape:
    tape.watch(x)			# 在梯度带上设置计算睡的梯度
    y = tf.pow(x, 4)		# 表达式y = x^4
    dy_dx = tape.gradient(y, x)	# 计算一届梯度后, y_ = 4*x^3
    print(dy_dx)				# 带入x = 3.0后, 得出梯度y_ = 108
del tape

2. 计算二阶梯度

x = tf.Variable(4.)
with tf.GradientTape() as tape1:
    tape1.watch(x)
    with tf.GradientTape() as tape2:
        tape2.watch(x)
        y = tf.pow(x, 4)
        dy_dx = tape2.gradient(y, x) # y' = 4x^3
        print(dy_dx)
    dy2_dx2 = tape1.gradient(dy_dx, x) # y'' = 12x^2
    print(dy2_dx2)
del tape1,tape2

3. 计算多个函数的梯度
   这里需要注意要在GradientType中加上persistent=True, 否则计算一次梯度后, 第二次就失效了

x = tf.Variable(3.)
with tf.GradientTape(persistent=True) as tape:
    tape.watch(x)
    y = tf.pow(x, 4)        # y = x^4
    z = tf.multiply(x, 3)   # z = 3*x
    dy_dx = tape.gradient(y, x)
    dz_dx = tape.gradient(z, x)
    print(dy_dx)
    print(dz_dx)
del tape

4. 最后在我们网络模型训练的时候, gradientTape会自动求解模型中可训练的变量的梯度,然后进行训练

with tf.GradientTape() as tape:
    pred = self.d.call(img, training=True)
    loss = self.loss_func(label, pred)
grads = tape.gradient(loss, self.d.trainable_variables)
self.opt.apply_gradients(zip(grads, self.d.trainable_variables))
return loss, binary_accuracy(label, pred)