Tensorflow: Как написать op с градиентом в python?

Да, как указано в ответе @ Yaroslav, это возможно, и ключ-ссылки, на которые он ссылается: здесь и здесь. Я хочу развить этот ответ, приведя конкретный пример.

Modulo opperation: давайте реализуем элементарную операцию по модулю в tensorflow (она уже существует, но ее градиент не определен, но для примера мы будем реализовывать ее с нуля).

Функция Numpy: первым шагом является определение opperation мы хотим для массивов numpy. Поэлементная модулярная оппозиция уже реализована в numpy, так что это легко:

import numpy as np
def np_mod(x,y):
    return (x % y).astype(np.float32)

Причина для .astype(np.float32) заключается в том, что по умолчанию tensorflow принимает типы float32, и если вы дадите ему float64 (по умолчанию numpy), он будет жаловаться.

Функция градиента: Далее нам нужно определить функцию градиента для нашей опперации для каждого входа опперации как функцию тензорного потока. Функция должна принимать очень специфическую форму. Он должен взять тензорное представление opperation op и градиент выходного сигнала grad и сказать, как распространять градиенты. В нашем случае градиенты опперации mod легки, производная равна 1 относительно первого аргумента и по отношению ко второму (почти везде и бесконечно при конечном числе точек, но давайте это проигнорируем, см. https://math.stackexchange.com/questions/1849280/derivative-of-remainder-function-wrt-denominator для деталей). Итак, у нас есть

def modgrad(op, grad):
    x = op.inputs[0] # the first argument (normally you need those to calculate the gradient, like the gradient of x^2 is 2x. )
    y = op.inputs[1] # the second argument

    return grad * 1, grad * tf.neg(tf.floordiv(x, y)) #the propagated gradient with respect to the first and second argument respectively

Функция grad должна возвращать N-кортеж, где n-число аргументов операции. Обратите внимание, что нам нужно вернуть функции тензорного потока входных данных.

Создание функции TF с градиентами: Как объясняется в источниках, упомянутых выше, существует хак для определения градиентов функции с помощью tf.RegisterGradient [doc] и tf.Graph.gradient_override_map [doc] .

Копируя код из harpone , мы можем модифицировать функцию tf.py_func, чтобы она одновременно определяла градиент:

import tensorflow as tf

def py_func(func, inp, Tout, stateful=True, name=None, grad=None):

    # Need to generate a unique name to avoid duplicates:
    rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))

    tf.RegisterGradient(rnd_name)(grad)  # see _MySquareGrad for grad example
    g = tf.get_default_graph()
    with g.gradient_override_map({"PyFunc": rnd_name}):
        return tf.py_func(func, inp, Tout, stateful=stateful, name=name)

Опция stateful должна сказать tensorflow, всегда ли функция дает один и тот же выход для одного и того же входа (stateful = False), и в этом случае tensorflow может просто построить график tensorflow, это наш случай и, вероятно, будет иметь место в большинстве ситуаций.

Комбинируя все это вместе: теперь, когда у нас есть все части, мы можем объединить их все вместе:

from tensorflow.python.framework import ops

def tf_mod(x,y, name=None):

    with ops.op_scope([x,y], name, "mod") as name:
        z = py_func(np_mod,
                        [x,y],
                        [tf.float32],
                        name=name,
                        grad=modgrad)  # <-- here's the call to the gradient
        return z[0]

tf.py_func действует на списки тензоров (и возвращает список тензоров), поэтому мы имеем [x,y] (и возвращаем z[0]). А теперь мы закончили. И мы можем это проверить.

Тест:

with tf.Session() as sess:

    x = tf.constant([0.3,0.7,1.2,1.7])
    y = tf.constant([0.2,0.5,1.0,2.9])
    z = tf_mod(x,y)
    gr = tf.gradients(z, [x,y])
    tf.initialize_all_variables().run()

    print(x.eval(), y.eval(),z.eval(), gr[0].eval(), gr[1].eval())

[ 0.30000001 0.69999999 1.20000005 1.70000005] [ 0.2 0.5 1. 2.9000001] [ 0.10000001 0.19999999 0.20000005 1.70000005] [ 1. 1. 1. 1.] [ -1. -1. -1. 0.]

Успех!