python - Tensorflow:gradient_override_map 无法覆盖 op tf.stack 的反向梯度

Question

我正在尝试使用 tf.RegisterGradient 和 tf.gradient_override_map 编辑 tf.stack op 的反向梯度计算机制，这是我的代码:

import tensorflow as tf

class SynthGradBuilder(object):
    def __init__(self):
        self.num_calls = 0

    def __call__(self, x, l=1.0):
        op_name = "SynthGrad%d" % self.num_calls
        @tf.RegisterGradient(op_name)
        def _grad_synth(op, grad):
            return grad[0]

        g = tf.get_default_graph()
        with g.gradient_override_map({"stack": op_name}):
            y = tf.stack([x,x])

        self.num_calls += 1
        return y

GradSys = SynthGradBuilder()

在另一个脚本中，我写了

import tensorflow as tf
from gradient_synthesizer import GradSys

x = tf.Variable([1,2])
y = GradSys(x, l=1)
z = tf.stack([x,x])


grad = tf.gradients(y, x, grad_ys=[[tf.convert_to_tensor([3, 4]), 
                              tf.convert_to_tensor([6, 8])]])
grad_stack = tf.gradients(z, x, grad_ys=[[tf.convert_to_tensor([3, 4]), 
                              tf.convert_to_tensor([6, 8])]])

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    print "grad bp: ", sess.run(grad)
    print "grad_stack: ", sess.run(grad_stack)
    print "y: ", sess.run(y)

预期的输出应该是:

grad bp: [3,4];
grad_stack: [3+6, 4+8] = [9, 12];
y: [[1,2], [1,2]];

我从代码中实际得到的是:

表明 tf.stack 的反向梯度根本没有被替换，这与我的预期相反。

我不确定这种差异是不是错误地使用“stack”作为操作tf.stack的类型串造成的，我做了如下实验:

描述张量y的第一项，“stack:0”暗示op tf.stack的注册名是“stack”，也是它的类型string。所以看起来这不是“堆栈”的错。

我无法找出代码问题的原因。我想知道是否有人可以帮助我。

Answer 1

Tl;dr:正确的代码应该是:

@tf.RegisterGradient(op_name)
def _grad_synth(op, grad):
  x, y = tf.unstack(grad)
  return [x, tf.zeros_like(y)]

g = tf.get_default_graph()
with g.gradient_override_map({"Pack": op_name}):
  y = tf.stack([x, x])

---

因为这是一个很常见的问题，所以我想解释一下更多细节:

您的原始代码中存在两个主要问题:

1. gradient_override_map 的错误用法:

tf.stack 的实际 OP 名称是 Pack(不是 Stack)，因此您需要覆盖 Pack 而不是 Stack:

`g.gradient_override_map({"Pack": op_name})`.

您可能想知道我怎么知道实际的 OP 名称？好吧，一个简单的方法是通过运行以下代码来探测 GraphDef:

with tf.Graph().as_default():
  x = tf.constant(0)
  y = tf.stack([x, x])
  print(tf.get_default_graph().as_graph_def())

1. 错误的梯度函数:

Pack 的原始渐变是一个简单的Unpack ( official code )。在你的情况下，你仍然需要先解压梯度，但只传播第一部分:

@tf.RegisterGradient(op_name)
def _grad_synth(op, grad):
  x, y = tf.unstack(grad)
  return [x, tf.zeros_like(y)]

请注意，此代码非常适合您的情况。然而，如果你想支持任意长度的堆栈，你可以使用稍微复杂一点的版本:

@tf.RegisterGradient(op_name)
def _grad_synth(op, grad):
  x_list = tf.unstack(grad)
  for i in range(1, len(x_list)):
    x_list[i] = tf.zeros_like(x_list[i])
  return x_list