TensorFlow：怎么替换或者修改梯度？

问题：

我想要去替换或修改 TensorFlow 中一个 op 或 部分计算图的梯度，该怎么整？

在某些情况下，这和 tf.stop_gradient() 是相对的：我们想要的不是“添加一个 op，然后在计算梯度的过程中忽略该 op”，我们想要的是“只在梯度计算时起作用（I want a calculation which is only used when calculating gradients）”

一个简单的例子如下：

在不改变前向传播的情况下，通过给梯度乘以一个常量来缩放梯度。

另一个例子如下：

在不改变前向传播的过程下，将梯度裁剪到给定的范围。

版本1
首先定义你自己的梯度计算函数，并注册到 tensorflow

@tf.RegisterGradient("CustomGrad") # 注册到tensorflow
def _const_mul_grad(unused_op, grad):
  return 5.0 * grad

因为我们不想改变前向传播过程，因此下面使用上面注册的函数替换原始的梯度计算函数（即用 “CustomGrad” 覆盖 “Identity”）

g = tf.get_default_graph()
with g.gradient_override_map({"Identity": "CustomGrad"}): # 覆盖该 op 的梯度计算函数
  output = tf.identity(input, name="Identity")

下面是一个梯度裁剪的实例：

import tensorflow as tf

@tf.RegisterGradient("CustomClipGrad")
def _clip_grad(unused_op, grad):
  return tf.clip_by_value(grad, -0.1, 0.1)

input = tf.Variable([3.0], dtype=tf.float32)

g = tf.get_default_graph()
with g.gradient_override_map({"Identity": "CustomClipGrad"}):
  output_clip = tf.identity(input, name="Identity")
grad_clip = tf.gradients(output_clip, input)

# output without gradient clipping in the backwards pass for comparison:
output = tf.identity(input)
grad = tf.gradients(output, input)

with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  print("with clipping:", sess.run(grad_clip)[0])
  print("without clipping:", sess.run(grad)[0])

版本2
在 TensorFlow 1.7 版本后，有一种定义梯度的快捷方式（也可以用来同时重定义多个op的梯度）

定义一个层，该层在保持正向传播不变的情况下，缩放梯度：

@tf.custom_gradient
def scale_grad_layer(x):
  def grad(dy):
    return 5.0 * dy
  return tf.identity(x), grad

下面是一个反向传播过程中梯度裁剪的例子：

import tensorflow as tf

input = tf.Variable([3.0], dtype=tf.float32)

@tf.custom_gradient
def clip_grad_layer(x):
  def grad(dy):
    return tf.clip_by_value(dy, -0.1, 0.1)
  return tf.identity(x), grad

output_clip = clip_grad_layer(input)
grad_clip = tf.gradients(output_clip, input)

# output without gradient clipping in the backwards pass for comparison:
output = tf.identity(input)
grad = tf.gradients(output, input)

with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  print("input:", sess.run(input))
  print("output_clipping:", sess.run(output_clip))
  print("output_no_clipping:", sess.run(output))
  print("with clipping:", sess.run(grad_clip)[0])
  print("without clipping:", sess.run(grad)[0])

使用 optimizer.compute_gradients 或 tf.gradient 来获得原始梯度，然后进行任何你需要的变换，最后使用 optimizer.apply_gradients

回答3

回答人：Bily
使用梯度截断，构造函数

假设前向传播过程如下：

y = f(x)

反向传播过程如下：

y = b(x)

一个简单的修改梯度的方法如下：

y = b(x) + tf.stop_gradient(f(x) - b(x))

实例：

def clip_gradient(x, l=1.0):
	positive_path = tf.stop_gradient(x * tf.cast(1 - l, tf.float32))
	negative_path = x * tf.cast(l, tf.float32)
	return positive_path + negative_path

需要注意的是，l也可以使用平滑，并不是一个常数，而是由0变为1，即

其中， [公式] 是一个超参数，文章中设为10； [公式] 随着训练的进行由0变为1，表示当前的训练步数/总的训练步数。上面的式子意味着一开始时， [公式] ，领域分类损失不会回传到编码器网络中，只有领域分类器得到训练；随着训练的进行， [公式] 逐渐增加，编码器得到训练，并开始逐步生成可以混淆领域分类器的特征。

部分参考：
https://blog.csdn.net/u014061630/article/details/81369787