Как сделать пользовательскую функцию активации только с Python в Tensorflow?

Question

Предположим, вам нужно создать функцию активации, которая невозможна с использованием только предварительно определенных стандартных блоков тензорного потока. Что вы можете сделать?

Таким образом, в Tensorflow можно сделать свою собственную функцию активации. Но это довольно сложно, вы должны написать это на C++ и перекомпилировать весь тензор потока [1] [2].

Есть ли более простой способ?

Answer 1

Да, есть!

Кредит: было трудно найти информацию и заставить ее работать, но вот пример, копирующий принципы и код, найденный здесь и здесь.

Требования: перед тем, как мы начнем, есть два требования, чтобы это могло быть успешным. Во-первых, вы должны быть в состоянии написать свою активацию как функцию для numy массивов. Во-вторых, вы должны иметь возможность записать производную этой функции либо в виде функции в Tensorflow (проще), либо в худшем случае в качестве функции для массивов с нулевыми значениями.

Запись Активация функции:

Итак, давайте возьмем для примера эту функцию, которую мы хотели бы использовать функцию активации:

def spiky(x):
    r = x % 1
    if r <= 0.5:
        return r
    else:
        return 0

Который выглядит следующим образом:

Первый шаг превращает его в функцию numpy, это легко:

import numpy as np
np_spiky = np.vectorize(spiky)

Теперь мы должны написать его производную.

Градиент активации: В нашем случае это просто, это 1, если x mod 1 < 0,5, и 0 в противном случае. Так:

def d_spiky(x):
    r = x % 1
    if r <= 0.5:
        return 1
    else:
        return 0
np_d_spiky = np.vectorize(d_spiky)

Теперь о сложной части создания функции TensorFlow.

Делаем ncty fct в тензор потока fct: Начнем с превращения np_d_spiky в функцию тензорного потока. В тензорном потоке есть функция tf.py_func(func, inp, Tout, stateful=stateful, name=name) [doc], который преобразует любую функцию numpy в функцию тензорного потока, поэтому мы можем использовать ее:

import tensorflow as tf
from tensorflow.python.framework import ops

np_d_spiky_32 = lambda x: np_d_spiky(x).astype(np.float32)


def tf_d_spiky(x,name=None):
    with tf.name_scope(name, "d_spiky", [x]) as name:
        y = tf.py_func(np_d_spiky_32,
                        [x],
                        [tf.float32],
                        name=name,
                        stateful=False)
        return y[0]

tf.py_func действует в списках тензоров (и возвращает список тензоров), поэтому мы имеем [x] (и вернуться y[0]). stateful опция заключается в том, чтобы сообщить tenorflow, всегда ли функция выдает один и тот же вывод для одного и того же ввода (stateful = False), и в этом случае тензор потока может просто отображать график тензорного потока, это наш случай и, вероятно, так будет в большинстве ситуаций. На данный момент нужно быть осторожным с тем, чтобы float64 но тензор потока использует float32 так что вам нужно преобразовать свою функцию для использования float32 прежде чем вы сможете преобразовать его в функцию tenorflow, иначе tenorflow будет жаловаться. Вот почему нам нужно сделать np_d_spiky_32 первый.

Как насчет градиентов? Проблема только с выполнением вышеизложенного состоит в том, что, хотя у нас теперь есть tf_d_spiky которая является тензорной версией np_d_spikyмы не могли бы использовать его как функцию активации, если бы захотели, потому что тензор потока не знает, как рассчитать градиенты этой функции.

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

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)

Теперь мы почти закончили, единственное, что функция grad, которую мы должны передать вышеупомянутой функции py_func, должна принять специальную форму. Он должен принимать в операции и предыдущие градиенты до операции и распространять градиенты в обратном направлении после операции.

Функция градиента: Итак, для нашей колючей функции активации мы так и сделаем:

def spikygrad(op, grad):
    x = op.inputs[0]

    n_gr = tf_d_spiky(x)
    return grad * n_gr  

Функция активации имеет только один вход, поэтому x = op.inputs[0], Если бы операция имела много входов, нам нужно было бы вернуть кортеж, один градиент для каждого входа. Например, если операция была a-bградиент по отношению к a является +1 и в отношении b является -1 так что мы бы return +1*grad,-1*grad, Обратите внимание на то, что нам нужно возвращать тензорные функции ввода, поэтому нужно tf_d_spiky, np_d_spiky не работал бы, потому что это не может воздействовать на тензоры тензора. В качестве альтернативы мы могли бы написать производную с использованием функций тензорного потока:

def spikygrad2(op, grad):
    x = op.inputs[0]
    r = tf.mod(x,1)
    n_gr = tf.to_float(tf.less_equal(r, 0.5))
    return grad * n_gr  

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

np_spiky_32 = lambda x: np_spiky(x).astype(np.float32)

def tf_spiky(x, name=None):

    with tf.name_scope(name, "spiky", [x]) as name:
        y = py_func(np_spiky_32,
                        [x],
                        [tf.float32],
                        name=name,
                        grad=spikygrad)  # <-- here's the call to the gradient
        return y[0]

И теперь мы закончили. И мы можем это проверить.

Тестовое задание:

with tf.Session() as sess:

    x = tf.constant([0.2,0.7,1.2,1.7])
    y = tf_spiky(x)
    tf.initialize_all_variables().run()

    print(x.eval(), y.eval(), tf.gradients(y, [x])[0].eval())

[0.2 0.69999999 1.20000005 1.70000005] [0.2 0. 0.200005 0.] [ 1. 0. 1. 0.]

Успех!

Answer 2

Почему бы просто не использовать функции, которые уже доступны в tenorflow, для создания новой функции?

Для spiky функция в вашем ответе, это может выглядеть следующим образом

def spiky(x):
    r = tf.floormod(x, tf.constant(1))
    cond = tf.less_equal(r, tf.constant(0.5))
    return tf.where(cond, r, tf.constant(0))

Я бы посчитал это существенно проще (даже не нужно вычислять какие-либо градиенты), и если вы не хотите делать действительно экзотические вещи, я едва ли могу представить, что тензорный поток не обеспечивает строительные блоки для создания очень сложных функций активации.