tensorflow自定义梯度 RegisterGradient gradient_override_map

tensorflow自定义梯度

1：梯度计算

关于tensorflow的反向传播算法可以查看我之前的文章反向传播(back propagation)算法详解，tensorflow可以自动计算网络结构的梯度，但是我们需要首先确保模型的梯度是能够正确学习的，而不是ill-posed的，比如说将softmax-loss里面的交叉熵换成一般的squared error的话，会存在梯度一直是0的现象，具体情况可以查看李宏毅的机器学习课程中关于回归以及分类的章节。

Tensorflow官网里面关于梯度有以下描述：

Given a graph of ops, TensorFlow uses automatic differentiation (backpropagation) to add new ops representing gradients with respect to the existing ops. To make automatic differentiation work for new ops, you must register a gradient function which computes gradients with respect to the ops' inputs given gradients with respect to the ops' outputs.

关于tensorflow 梯度部分的内容可以查看https://www.tensorflow.org/guide/extend/op#implement_the_gradient_in_python。

2：RegisterGradient

我这里说一下自定义梯度问题，在domain-adversial算法(transfer learning)里面，需要adversial部分的梯度对原来的net起的是副作用：

因为对于transfer learning来说，我们希望学习到的feature处于同一分布，那么这一个网络就能够应用到其他的输入上，这就是domain-adversial的思想。vgg是常用的提取图像特征的网络，基于imagenet训练好的vgg网络能够直接用到图像特征提取上面，不需要重新训练，这就是transfer learning。

domain-adversial使用了一个domain classifier，目的就是分清楚输入X时来自于哪个。这里说明一下，就是比如说使用狗和猫的图片训练了一个分类网络，判断输入是猫还是狗，但是现在我们要应用到大象和狮子上，用于训练的数据和需要transfer的数据差异较大，因此直接使用的话，结果会很差。

因此，我们希望绿色网络部分得到的猫狗狮子大象的特征处于同一分布，那么训练好的网络就可以直接应用到别的domain下面。domain classifier的思想和GAN类似，domain classifier希望吧不同分布的数据区分开，但是绿色的网络不希望clssifier分开（希望学习到的特征处于同一分布，所以很难分开），因此需要一个reversal gradient过程，因为两个网络的目的不一样，classifier希望自己尽可能准确，feature extractor希望产生的特征使得classifer尽可能不正确。

类似的思想也体现在了acm multimedia 2017 best paper上面：

下面的代码实现的是clip gradient:

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])

运行一下可以看到，原本identity连接，梯度应该是1，但是结果变成了0.1.

下面的代码将梯度减半：

import tensorflow as tf

@tf.RegisterGradient("A")
def _clip_grad(unused_op, grad): # grad为当前tensorflow中的默认真实梯度
    return [grad*0.5]

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

g = tf.get_default_graph()
with g.gradient_override_map({"Identity": "A"}):
    # 当x=3时，正向传播（不受影响）: y  = (6x)^2 = (6x3)^2 = 324, 如下：
    # 当x=3时，反向传播（人为定义）：y' = 72x = 72*3 = 216, 如下：
    # 因为我们重新定义了反向传播梯度的计算方式，因此结果是108
    output_clip = tf.square(tf.multiply(Input,6,name="Identity"))
    output_A=tf.identity(output_clip)
    
grad_clip = tf.gradients(output_A, 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())
    o = sess.run(output_clip)
    print("forward prop: output = ",o)
    print("with clipping:", sess.run(grad_clip))
    print("without clipping:", sess.run(grad))

这些例子都比较简单，一些复杂的例子可以查看https://stackoverflow.com/questions/43839431/tensorflow-how-to-replace-or-modify-gradient/43948872#43948872。

另外，domain-adversial的reversal gradient通过一下代码实现：

import tensorflow as tf
from tensorflow.python.framework import ops

#inverse gradient in BP
class FlipGradientBuilder(object):
    def __init__(self):
        self.num_calls = 0
    #callable: execute follwing ops when call FlipGradientBuilder(obj)
    def __call__(self, x, l=1.0):
        grad_name = "FlipGradient%d" % self.num_calls
        @ops.RegisterGradient(grad_name)
        def _flip_gradients(op, grad):
            return [tf.negative(grad) * l]

        g = tf.get_default_graph()
        with g.gradient_override_map({"Identity": grad_name}):
            y = tf.identity(x)

        self.num_calls += 1
        return y

flip_gradient = FlipGradientBuilder()

这些例子都是通过registerGradient实现的，关于这部分，可以产看https://tensorflow.google.cn/api_docs/python/tf/Graph#gradient_override_map。

有一个常见的疑问就是为什么这些例子中都是讲tf.identity转化为自定义的ops，原因就是如果直接是这样比较好懂一点，不仅仅是identity,所有的tenorflow op比如说matmul之类的操作都可以被重写反向传播梯度。