tensorflow - 选择满足特定条件的tensorflow中的索引

Question

我希望选择矩阵中元素的坐标满足特定条件的矩阵元素。例如，条件可以是:(y_coordinate-x_coordinate) == -4因此，将选择那些坐标满足此条件的元素。如何在不遍历每个元素的情况下有效地做到这一点？

Answer 1

也许你需要tf.gather_nd :

iterSession = tf.InteractiveSession()

vals = tf.constant([[1,2,3], [4,5,6], [7,8,9]])
arr = tf.constant([[x, y] for x in range(3) for y in range(3) if -1 <= x - y <= 1])

arr.eval()
# >> array([[0, 0],
# >>        [0, 1],
# >>        [1, 0],
# >>        [1, 1],
# >>        [1, 2],
# >>        [2, 1],
# >>        [2, 2]], dtype=int32)

tf.gather_nd(vals, arr).eval()
# >> array([1, 2, 4, 5, 6, 8, 9], dtype=int32)

---

或tf.boolean_mask :

iterSession = tf.InteractiveSession()

vals = tf.constant([[1,2,3], [4,5,6], [7,8,9]])
arr = tf.constant([[-1 <= x - y <= 1 for x in range(3)] for y in range(3)])

arr.eval()
# array([[ True,  True, False],
#        [ True,  True,  True],
#        [False,  True,  True]], dtype=bool)

tf.boolean_mask(vals, arr).eval()
# array([ 1.,  2.,  4.,  5.,  6.,  8.,  9.], dtype=int32)