Tensorflow自定义激活函数/函数/梯度

Tensorflow自定义激活函数/函数/梯度

(对于Tensorflow 1.x)
最近刚做完一个paper，需要自定义激活函数，记录一下心得，顺便帮助下有需要的小伙伴。大刀阔斧，直接上解决方案：

1、对于分段（激活）函数，代码分开写

2、使用自带自定义梯度

详解

Tensorflow是自动求导（不懂就百度），因此我们不需要定义梯度，但大家可能会遇到和我一样的问题（在训练模型的时候loss爆炸），所以大家才会来查吧。
自定义激活函数/函数直接定义就可以，比如：

output = tf.exp(input)
output = tf.log(input)

但为什么有时候会梯度爆炸？
因为激活函数大多是参照relu进行修改，故大多是分段函数，分段函数在tensorflow中使用

tf.where(tf.greater(input, [0.0]),function1,function2)

funtion1计算大于0的数，function2计算小于等于0的数，但这就导致我构造的激活函数loss爆炸。原因不详，猜测是先计算所有输入都参与function1和function2的计算。
我使用了tensorflow定义swish的例子定义函数：

def _swish_shape(op):
  return [op.inputs[0].shape]

@function.Defun(shape_func=_swish_shape, func_name="swish_grad", noinline=True)
def _swish_grad(features, grad):
  sigmoid_features = math_ops.sigmoid(features)
  activation_grad = (
      sigmoid_features * (1.0 + features * (1.0 - sigmoid_features)))
  return grad * activation_grad

@tf_export("nn.swish")
@function.Defun(
    grad_func=_swish_grad,
    shape_func=_swish_shape,
    func_name="swish",
    noinline=True)
def swish(features):
  features = ops.convert_to_tensor(features, name="features")
  return features * math_ops.sigmoid(features)

加入可训练参数的话，因为涉及tensorflow自带函数，较为复杂，可以用上述形式，定义多个函数实现。
下面是一个简单例子

#F2
@function.Defun(shape_func=shape, func_name="F2_grad", noinline=True)
def F2_grad(features, grad):
    return grad * tf.where(tf.greater(features, features*0), features*0+1, tf.exp(features))


@tf_export("F2")
@function.Defun(
    grad_func=F2_grad,
    shape_func=shape,
    func_name="F2",
    noinline=True
)
def F2(features):
    return tf.where(tf.greater(features, features*0), features, tf.exp(features))




def my_best(features, a1, a2, a3):
    f1 = F1(a2*features)
    f2 = F2(a3*features)
    return tf.where(tf.greater(features, features*0), features, a1*(f1-features*f2))