如何为具有多个输入的TensorFlow op分配自定义渐变

时间:2019-01-04 23:45:10

标签: python tensorflow tensorflow-gradient

我正在尝试使用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)这样的元组中,以使该函数只有“两个”输出。 TensorFlow也拒绝了这一点,这是因为事后看来,它不能像Tensor那样调用元组-完全合理。

我已经对示例进行了更多讨论,但是没有用。无论我尝试什么,都无法获得自定义的渐变来处理多个输入。我希望比我有更多关于自定义操作和渐变的知识的人对此会有更好的主意-预先感谢您的帮助!

2 个答案:

答案 0 :(得分:1)

前段时间,我遇到了类似的问题,我认为文档对此不太清楚。通常,代码应类似于:

@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函数应返回N次梯度,其中N是custom_operation用作输入的参数数目(除范围外)。通过使用两个输入(x和y),grad函数必须返回两次梯度(对于x一次,对于y一次)。通常,对于两个输入,您还可以使grad()函数返回g1!= g2而不是g。 因此,在您的示例中,它变为:

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

    def grad(dy):
        return dy, dy

    return tf.identity(x*z), grad

答案 1 :(得分:0)

如果我们使用多个变量作为输入,则从“ grad”函数返回的渐变数量应等于输入变量的数量,尽管我们可能并不在意其中的一些。

例如:

@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”的第二个输出是一个函数,而不是梯度张量。