Tensorflow：如何替換或修改漸變？

Question

我想用張量流來替換或修改一個op或部分圖的梯度。如果我可以在計算中使用現有的漸變，那將是理想的。

在某些方面，這與tf.stop_gradient()所做的相反：不是在計算梯度時添加被忽略的計算，而是需要一個僅在計算梯度時使用的計算。

一個簡單的例子就是通過將梯度乘以常數來簡單地縮放梯度（但不會將正向計算乘以常數）。另一個例子是將漸變剪輯到給定範圍的東西。

Answer 1

首先定義自定義梯度：

@tf.RegisterGradient("CustomGrad") 
def _const_mul_grad(unused_op, grad): 
    return 5.0 * grad 

既然你什麼都不想要在直傳發生的，跟你的新梯度覆蓋的身份運行的梯度：

g = tf.get_default_graph() 
with g.gradient_override_map({"Identity": "CustomGrad"}): 
    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]) 

Answer 2

使用optimizer.compute_gradients或tf.gradient獲得原始梯度
然後做你想做的
最後，使用optimizer.apply_gradients

我發現了一個example從GitHub

Answer 3

做到這一點的最一般的方法是使用 https://www.tensorflow.org/api_docs/python/tf/RegisterGradient

下面我實現了反向傳播的漸變剪裁，它可以用於matmul，如這裏所示，或者任何其它運算：

import tensorflow as tf 
import numpy as np 

# from https://gist.github.com/harpone/3453185b41d8d985356cbe5e57d67342 
def py_func(func, inp, Tout, stateful=True, name=None, grad=None): 

    # Need to generate a unique name to avoid duplicates: 
    rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8)) 

    tf.RegisterGradient(rnd_name)(grad) 
    g = tf.get_default_graph() 
    with g.gradient_override_map({"PyFunc": rnd_name}): 
     return tf.py_func(func, inp, Tout, stateful=stateful, name=name) 

def clip_grad(x, clip_value, name=None): 
    """" 
    scales backpropagated gradient so that 
    its L2 norm is no more than `clip_value` 
    """ 
    with tf.name_scope(name, "ClipGrad", [x]) as name: 
     return py_func(lambda x : x, 
         [x], 
         [tf.float32], 
         name=name, 
         grad=lambda op, g : tf.clip_by_norm(g, clip_value))[0] 

---

實例：

with tf.Session() as sess: 
    x = tf.constant([[1., 2.], [3., 4.]]) 
    y = tf.constant([[1., 2.], [3., 4.]]) 

    print('without clipping') 
    z = tf.matmul(x, y) 
    print(tf.gradients(tf.reduce_sum(z), x)[0].eval()) 

    print('with clipping') 
    z = tf.matmul(clip_grad(x, 1.0), clip_grad(y, 0.5)) 
    print(tf.gradients(tf.reduce_sum(z), x)[0].eval()) 

    print('with clipping between matmuls') 
    z = tf.matmul(clip_grad(tf.matmul(x, y), 1.0), y) 
    print(tf.gradients(tf.reduce_sum(z), x)[0].eval()) 

---

輸出：

without clipping 
[[ 3. 7.] 
[ 3. 7.]] 
with clipping 
[[ 0.278543 0.6499337] 
[ 0.278543 0.6499337]] 
with clipping between matmuls 
[[ 1.57841039 3.43536377] 
[ 1.57841039 3.43536377]] 

Answer 4

假設向前計算是

y = f(x) 

而且你希望它像backpropagate

y = b(x) 

一個簡單的黑客將是：

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