Tensorflow 学习笔记 -----tf.where

TensorFlow函数：tf.where

在之前版本对应函数tf.select

官方解释:

 

 1 tf.where(input, name=None)`
 2 Returns locations of true values in a boolean tensor.
 3 
 4 This operation returns the coordinates of true elements in input. The coordinates are returned in a 2-D tensor where the first dimension (rows) represents the number of true elements, and the second dimension (columns) represents the coordinates of the true elements. Keep in mind, the shape of the output tensor can vary depending on how many true values there are in input. Indices are output in row-major order.
 5 
 6 For example:
 7 # 'input' tensor is [[True, False]
 8 #                    [True, False]]
 9 # 'input' has two true values, so output has two coordinates.
10 # 'input' has rank of 2, so coordinates have two indices.
11 where(input) ==> [[0, 0],
12                   [1, 0]]
13 
14 # `input` tensor is [[[True, False]
15 #                     [True, False]]
16 #                    [[False, True]
17 #                     [False, True]]
18 #                    [[False, False]
19 #                     [False, True]]]
20 # 'input' has 5 true values, so output has 5 coordinates.
21 # 'input' has rank of 3, so coordinates have three indices.
22 where(input) ==> [[0, 0, 0],
23                   [0, 1, 0],
24                   [1, 0, 1],
25                   [1, 1, 1],
26                   [2, 1, 1]]

 

有两种用法：

1、tf.where(tensor)

tensor 为一个bool 型张量，where函数将返回其中为true的元素的索引。如上图官方注释

2、tf.where(tensor,a,b)

a,b为和ｔｅｎｓｏｒ相同维度的ｔｅｎｓｏｒ，将ｔｅｎｓｏｒ中的ｔｒｕｅ位置元素替换为ａ中对应位置元素，ｆａｌｓｅ的替换为ｂ中对应位置元素。

例：

import tensorflow as tf
import numpy as np
sess=tf.Session()

a=np.array([[1,0,0],[0,1,1]])
a1=np.array([[3,2,3],[4,5,6]])

print(sess.run(tf.equal(a,1)))
print(sess.run(tf.where(tf.equal(a,1),a1,1-a1)))

 

＞＞［［ｔｒｕｅ，ｆａｌｓｅ，ｆａｌｓｅ］，［ｆａｌｓｅ，ｔｒｕｅ，ｔｒｕｅ］］

＞＞［［３，-１，-２］，［-３，５，６］］