Tensorflow 다른 자리 표시 자에 따라 자리 표시 자에서 요소 가져 오기

Question

tf에서이 작업을 코딩하는 데 문제가 있습니다. 다음은 예를 들어 하나의 [n, 2] 자리 표시 자 x와 하나의 [n, 1] 자리 표시 자 y가 있다고 가정합니다. x = [[1,2], [3,4], [5,6]] y = [1,0,1] y에서 각 요소 i에 대해 해당 요소를 i 번째 2d에서 가져 오려고합니다. 텐서. 예제에서 출력은 [2,3,6]이어야합니다. 몇 가지 기법을 시도했지만 성공하지 못했습니다. tensorflow로 쉽게 할 수 있습니까?

당신에게

Answer 1

사용 중 tf.gather_nd 감사 또는 tf.stack 및 tf.where하여 수동으로 해킹 :

import tensorflow as tf 

x = tf.convert_to_tensor([[1, 2], [3, 4], [5, 6]]) 
y = tf.convert_to_tensor([1, 0, 1]) 

with tf.Session() as sess: 
    xx = tf.unstack(x, axis=1) 
    ans = tf.where(tf.equal(y, tf.zeros_like(y)), xx[0], xx[1]) 
    print sess.run(ans) 


with tf.Session() as sess: 
    idx = tf.range(0, limit=3, delta=1, name='arange') 
    idx = tf.stack([idx, y], axis=-1) 
    ans = tf.gather_nd(x, idx) 
    print sess.run(ans)