tensorflow 的tf.where詳解

最近在用到數據篩選，觀看代碼中有tf.where()的用法，不是很常用，也不是很好理解。在這里記錄一下

tf.where(
    condition,
    x=None,
    y=None,
    name=None
)

Return the elements, either from x or y, depending on the condition.

理解：where嘛，就是要根據條件找到你要的東西。

condition：條件，是一個boolean

x：數據

y：同x維度的數據。

返回，返回符合條件的數據。當條件為真，取x對應的數據；當條件為假，取y對應的數據

舉例子。

def test_where():
    # 定義一個tensor，表示condition，內部數據隨機產生
    condition = tf.convert_to_tensor(np.random.random([5]), dtype=tf.float32)

    # 定義兩個tensor，表示原數據
    a = tf.ones(shape=[5, 3], name='a')

    b = tf.zeros(shape=[5, 3], name='b')

    # 選擇大於0.5的數值的坐標，並根據condition信息在a和b中選取數據
    result = tf.where(condition > 0.5, a, b)

    with tf.Session() as sess:
        print("condition:\n", sess.run([condition, result]))

結果：