如何在 Tensorflow 中制作一个只有 Python 的自定义激活函数?

How to make a custom activation function with only Python in Tensorflow?

假设您需要制作一个仅使用预定义的 tensorflow 构建块无法实现的激活函数,您可以做什么?

所以在 Tensorflow 中可以制作自己的激活函数。但是比较复杂,必须用C++写,然后重新编译整个tensorflow [1] [2].

有没有更简单的方法?

有!

来源: 很难找到信息并使其正常工作,但这里有一个从找到的原则和代码中复制的示例 here and here

要求: 在我们开始之前,有两个要求才能成功。首先,您需要能够将激活编写为 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 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 作用于张量列表(和 return 是张量列表),这就是为什么我们有 [x](和 return y[0] ). stateful 选项是告诉 tensorflow 该函数是否总是为相同的输入提供相同的输出(stateful = False),在这种情况下,tensorflow 可以简单地使用 tensorflow 图,这就是我们的情况,并且在大多数情况下可能都是这种情况情况。此时要注意的一件事是 numpy 使用 float64 但 tensorflow 使用 float32 因此您需要将函数转换为使用 float32 才能将其转换为 tensorflow 函数否则tensorflow 会抱怨。这就是为什么我们需要先制作np_d_spiky_32

梯度如何? 仅执行上述操作的问题是,即使我们现在有 tf_d_spiky,它是 np_d_spiky 的 tensorflow 版本,我们不能将它用作激活函数,因为 tensorflow 不知道如何计算该函数的梯度。

Hack to get Gradients: 正如上面提到的资源中所解释的,有一个 hack 可以使用 tf.RegisterGradient [doc] and tf.Graph.gradient_override_map [doc]. Copying the code from harpone 来定义函数的梯度可以 mod 化 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] 的原因。如果操作有很多输入,我们将需要 return 一个元组,每个输入一个梯度。例如,如果操作是 a-b 相对于 a 的梯度是 +1 而相对于 b-1 所以我们将有 return +1*grad,-1*grad.请注意,我们需要 return 输入的 tensorflow 函数,这就是为什么需要 tf_d_spikynp_d_spiky 不会起作用的原因,因为它不能作用于 tensorflow 张量。或者,我们可以使用张量流函数编写导数:

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.]

成功!

为什么不简单地使用 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 不会提供用于构建高度复杂的激活函数的构建块。