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

Question

Я пытаюсь использовать функциональность TensorFlow @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

Я столкнулся с подобной проблемой некоторое время назад, и я думаю, что документация не очень ясна по этому вопросу. В общем случае код должен выглядеть примерно так:

@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

Answer 2

Если мы используем несколько переменных в качестве входных данных, количество градиентов, возвращаемых из функции "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" - это функция, а не тензор градиента.