tf.where()详解

函数介绍

tf.where()返回一个布尔张量中真值的位置。对于非布尔型张量，非0的元素都判为True
返回的是二维张量，第一个维度的数量，即行数表明有多少个为True值；同一行中的数据代表该True值在原张量中的位置。
举几个例子：

test_a=tf.constant([[1,2,3],[0,0,6]])
tf.where(test_a)

输出：

test_a=tf.constant([[1,-2,-5],[0,0,6]])
tf.where(test_a)

输出：

test_a=tf.constant([[1,0,-5],[0,-1,0]])
tf.where(test_a)

输出为：

tf.where()还有另外一个用法：tf.where(input_tensor, a,b)
它的作用是：对于张量a，如果input_tensor对应位置的元素为True，则张量a中的该位置处元素保留，反之由张量b中相应位置的元素来代替。
这个用法可以用来自定义huber loss：

def huber_fn(y_true, y_pred):
    error = y_true - y_pred
    is_small_error = tf.abs(error) < 1
    squared_loss = tf.square(error) / 2
    linear_loss  = tf.abs(error) - 0.5
    return tf.where(is_small_error, squared_loss, linear_loss)