如何将自定义梯度分配给具有多个输入的 TensorFlow 操作

Question

我正在尝试使用 TensorFlow 的 @tf.custom_gradient 功能将自定义梯度分配给具有多个输入的函数。我可以只为一个输入组合一个工作设置，但不能为两个或更多输入组合。

我的代码基于 TensorFlow's custom_gradient documentation，它适用于一个输入，如本例所示：

import tensorflow as tf
import os

# Suppress Tensorflow startup info
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'

# Custom gradient decorator on a function,
# as described in documentation
@tf.custom_gradient
def my_identity(x):

    # The custom gradient
    def grad(dy):
        return dy

    # Return the result AND the gradient
    return tf.identity(x), grad

# Make a variable, run it through the custom op
x = tf.get_variable('x', initializer=1.)
y = my_identity(x)

# Calculate loss, make an optimizer, train the variable
loss = tf.abs(y)
opt = tf.train.GradientDescentOptimizer(learning_rate=0.001)
train = opt.minimize(loss)

# Start a TensorFlow session, initialize variables, train
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(train)

此示例静默运行，然后关闭。没有问题，没有错误。变量按预期优化。但是，在我的应用程序中，我需要对多个输入进行这样的计算，因此采用以下形式：

@tf.custom_gradient
def my_identity(x, z):

    def grad(dy):
        return dy

    return tf.identity(x*z), grad

运行代替示例（并向 my_identify 的调用添加另一个变量输入）导致以下错误输出。据我所知，错误的最后部分来自 op 的动态生成——信息格式与 op 建立所需的 C++ 格式相匹配（尽管这就是我所知道的全部）。

Traceback (most recent call last):
  File "testing.py", line 27, in <module>
    train = opt.minimize(loss)
  File "/usr/lib/python3/dist-packages/tensorflow/python/training/optimizer.py", line 400, in minimize
    grad_loss=grad_loss)
  File "/usr/lib/python3/dist-packages/tensorflow/python/training/optimizer.py", line 519, in compute_gradients
    colocate_gradients_with_ops=colocate_gradients_with_ops)
  File "/usr/lib/python3/dist-packages/tensorflow/python/ops/gradients_impl.py", line 630, in gradients
    gate_gradients, aggregation_method, stop_gradients)
  File "/usr/lib/python3/dist-packages/tensorflow/python/ops/gradients_impl.py", line 821, in _GradientsHelper
    _VerifyGeneratedGradients(in_grads, op)
  File "/usr/lib/python3/dist-packages/tensorflow/python/ops/gradients_impl.py", line 323, in _VerifyGeneratedGradients
    "inputs %d" % (len(grads), op.node_def, len(op.inputs)))
ValueError: Num gradients 2 generated for op name: "IdentityN"
op: "IdentityN"
input: "Identity"
input: "x/read"
input: "y/read"
attr {
  key: "T"
  value {
    list {
      type: DT_FLOAT
      type: DT_FLOAT
      type: DT_FLOAT
    }
  }
}
attr {
  key: "_gradient_op_type"
  value {
    s: "CustomGradient-9"
  }
}
 do not match num inputs 3

根据其他自定义渐变选项，我推测问题是缺少为第二个输入参数提供的渐变。所以，我将函数更改为：

@tf.custom_gradient
def my_identity(x, z):

    def grad(dy):
        return dy

    return tf.identity(x*z), grad, grad

这会导致以下更常见的错误：

Traceback (most recent call last):
  File "testing.py", line 22, in <module>
    y = my_identity(x, z)
  File "/usr/lib/python3/dist-packages/tensorflow/python/ops/custom_gradient.py", line 111, in decorated
    return _graph_mode_decorator(f, *args, **kwargs)
  File "/usr/lib/python3/dist-packages/tensorflow/python/ops/custom_gradient.py", line 132, in _graph_mode_decorator
    result, grad_fn = f(*args)
ValueError: too many values to unpack (expected 2)

@custom_gradient 装饰器仅将最后返回的元素标识为渐变。因此，我尝试将两个梯度作为 (grad, grad) 放入一个元组中，这样该函数将只有 "two" 输出。 TensorFlow 也拒绝了这一点，这次是因为它不能像调用 Tensor 那样调用元组——事后看来完全合理。

我又对这个例子大惊小怪，但无济于事。无论我尝试什么，我都无法获得自定义渐变来处理多个输入。我希望在自定义操作和渐变方面知识比我多的人对此有更好的想法——在此先感谢您的帮助！

Answer 1

如果我们使用多个变量作为输入，"grad" 函数的梯度数 return 应该等于输入变量的数量，尽管我们可能不关心其中的一些变量。

例如：

@tf.custom_gradient
def my_multiple(x,z):

def grad(dy):
    # return two gradients, one for 'x' and one for 'z'
    return (dy*z, dy*x)

return tf.identity(x*z), grad

注意"my_multiple"的第二个输出是函数，不是梯度张量。

Answer 2

我运行前段时间遇到过类似的问题，我认为文档对此不是很清楚。一般来说，代码应该是这样的：

@tf.custom_gradient
def custom_operation(x, y, scope='custom_op'):

    # define the gradient
    def grad(g):
        return g, g

    # define the forward pass (a multiplication, in this example)
    with tf.variable_scope(scope):
        forward_pass = x * y

    return forward_pass, grad

实际上，您的内部 grad 函数应该 return 梯度 N 次，其中 N 是 custom_operation 作为输入（除了范围）的参数数量。通过使用两个输入（x 和 y），grad 函数必须 return 两次梯度（一次用于 x，一次用于 y）。通常，您还可以为两个输入创建 grad() 函数 return g1 != g2 而不是 g。因此，在您的示例中，它变为：

@tf.custom_gradient
def my_identity(x, z):

    def grad(dy):
        return dy, dy

    return tf.identity(x*z), grad

如何将自定义梯度分配给具有多个输入的 TensorFlow 操作

How to assign custom gradient to TensorFlow op with multiple inputs

python

tensorflow

tensorflow-gradient