Как назначить пользовательский градиент для операции TensorFlow с несколькими входами

Question

Я пытаюсь использовать TensorFlow's @tf.custom_gradient функциональность для назначения собственного градиента функции с несколькими входами. Я могу собрать рабочую настройку только для одного входа, но не для двух или более.

Я основал свой код на документации custom_gradient TensorFlow, которая прекрасно работает для одного входа, как в этом примере:

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

Выполнение этого вместо примера (и добавление другой переменной input к вызову my_identify) приводит к следующей ошибке вывода. Насколько я могу судить, последние части ошибки связаны с динамическим генерированием операции - формат информации соответствует форматированию 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 - вполне разумно, оглядываясь назад.

Я возился с примером еще немного, но безрезультатно. Независимо от того, что я пытаюсь, я не могу получить пользовательский градиент для работы с несколькими входами. Я надеюсь, что кто-то с большим знанием, чем я, относительно пользовательских операций и градиентов, получит лучшее представление об этом - заранее спасибо за помощь!

Answer 1

Если мы используем несколько переменных в качестве входных данных, количество градиентов, возвращаемых из функции "grad", должно быть равно количеству входных переменных, хотя нам, возможно, нас это не волнует.

Например:

@tf.custom_gradient
def my_identity(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_identity" - это функция, а не тензор градиента.

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 должна возвращать градиент 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