如何在 Tensorflow 中仅使用 Python 创建自定义激活函数？

问题描述：

假设您需要创建一个激活函数，而这仅使用预定义的 TensorFlow 构建块是不可能实现的，您该怎么办？

因此，在 Tensorflow 中可以创建自己的激活函数。但这相当复杂，你必须用 C++ 编写它并重新编译整个 Tensorflow [1] [2]。

有没有更简单的方法？

解决方案 1：

有！

信用：
很难找到信息并使其发挥作用，但这里有一个从这里和这里的原理和代码中复制的示例。

要求：
在开始之前，有两个要求才能成功。首先，您需要能够将激活写为 numpy 数组上的函数。其次，您必须能够将该函数的导数写为 Tensorflow 中的函数（更简单），或者在最坏的情况下写为 numpy 数组上的函数。

编写激活函数：

因此，让我们以这个我们想要使用激活函数的函数为例：

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)

现在我们应该写出它的导数。

激活梯度：
在我们的例子中很简单，如果 x mod 1 < 0.5，则为 1，否则为 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 函数的困难部分。

将 numpy fct 转换为 tensorflow fct：我们首先将 np_d_spiky 转换为 tensorflow 函数。tensorflow tf.py_func(func, inp, Tout, stateful=stateful, name=name) [doc]
中有一个函数可以将任何 numpy 函数转换为 tensorflow 函数，因此我们可以使用它：

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选项是告诉 tensorflow 函数是否始终为相同的输入提供相同的输出（stateful = False），在这种情况下，tensorflow 可以简单地使用 tensorflow 图，这是我们的情况，并且在大多数情况下可能都是这种情况。此时要注意的一件事是 numpy 使用float64但 tensorflow 使用，float32因此您需要先将您的函数转换为使用float32，然后才能将其转换为 tensorflow 函数，否则 tensorflow 会抱怨。这就是我们需要np_d_spiky_32先制作的原因。

那么梯度呢？仅执行上述操作的问题是，即使我们现在有了tf_d_spiky的 TensorFlow 版本np_d_spiky，我们也无法将其用作激活函数，因为 TensorFlow 不知道如何计算该函数的梯度。

获取梯度的技巧：如上文提到的资料所述，有一种使用tf.RegisterGradient [doc]和tf.Graph.gradient_override_map [doc]定义函数梯度的技巧。复制harpone的代码，我们可以修改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)

现在我们几乎完成了，唯一需要做的是，我们需要传递给上述 py_func 函数的 grad 函数需要采用特殊形式。它需要接受一个操作，以及操作之前的先前梯度，并在操作之后向后传播梯度。

梯度函数：因此，对于我们的尖峰激活函数，我们将这样做：

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

    n_gr = tf_d_spiky(x)
    return grad * n_gr  

激活函数只有一个输入，这就是为什么x = op.inputs[0]。如果操作有许多输入，我们需要返回一个元组，每个输入一个梯度。例如，如果操作是关于和 的梯度a-b，那么我们就会有。请注意，我们需要返回输入的 TensorFlow 函数，这就是为什么需要，因为它不能作用于 TensorFlow 张量，所以它不起作用。或者，我们可以使用 TensorFlow 函数编写导数：a`+1b-1return +1*grad,-1*gradtf_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.20000005 0.] [ 1.0.1.0.]

成功！

解决方案 2：

为什么不简单地使用 TensorFlow 中已有的函数来构建新函数呢？

对于您答案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))

我认为这要容易得多（甚至不需要计算任何梯度），除非你想做一些非常奇特的事情，否则我几乎无法想象 TensorFlow 没有提供构建高度复杂的激活函数的基础模块。