如何在TensorFlow 2中实现含条件判断的自定义激活函数？

你需要的激活函数本质是对输入张量的每个元素执行符号保留的平方根运算，由于TensorFlow的张量是批量数据，不能用Python原生的if-else逐元素判断，必须使用TensorFlow提供的向量化操作来实现。以下是两种高效的实现方式：

方法一：使用tf.where显式分支判断 #

tf.where可以根据布尔条件张量，对输入张量的每个元素选择对应分支的计算结果，完全适配批量张量操作：

import tensorflow as tf

def sqrt_activation(x):
    # 生成布尔条件张量：每个元素判断是否 >= 0
    condition = tf.greater_equal(x, 0.0)
    # 正分支：直接计算平方根
    positive_output = tf.math.sqrt(x)
    # 负分支：先取反再算平方根，最后加负号保留原符号
    negative_output = -tf.math.sqrt(-x)
    # 根据条件选择对应输出
    return tf.where(condition, positive_output, negative_output)

方法二：利用数学等价式简化实现 #

观察你的函数逻辑可以发现，它等价于输入的符号乘以输入绝对值的平方根，用TensorFlow内置函数可以更简洁地实现：

import tensorflow as tf

def sqrt_activation(x):
    return tf.sign(x) * tf.sqrt(tf.abs(x))

两种实现都完全兼容TensorFlow 2的计算图机制，支持自动求导，且能高效处理任意形状的输入张量。

内容的提问来源于stack exchange，提问作者David H. J.