tensorflow根据tensor筛选指定行:boolean_mask

使用tensorflow时经常会遇到这样的一个问题，根据条件筛选留下指定的行，例如考虑tensor:

tensor_a = tf.Variable([[1,2,1],[2,3,0],[2,1,1],[2,5,1],[3,5,0]])

如果我们希望根据第3个特征即第3列大于0或小于0来筛选的话，应该如何实现呢？tensorflow提供了接口函数boolean_mask(x,y)，该函数会根据第2个tensor y的值来筛选x的行，y为boolean类型的矩阵，直接给一个例子：

import tensorflow as tf
tensor_a = tf.Variable([[1,2,1],[2,3,0],[2,1,1],[2,5,1],[3,5,0]])
tensor_b = tf.greater(tensor_a[:,2], 0)
tnesor_c = tf.boolean_mask(tensor_a, tensor_b)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print('*************')
    print(sess.run(tensor_b))
    print('*************')
    print(sess.run(tnesor_c))

结果如下：

*************
[ True False  True  True False]
*************
[[1 2 1]
 [2 1 1]
 [2 5 1]]