2017-08-06 165 views
0

我有一個data張量dimensios [B X N X 3],並且我有一個張量的維度[B X M]。我希望使用indices張量從data張量中提取[B X M X 3]張量。如何使用張量來索引tensorflow中的另一個張量

我有這樣的代碼工作:

new_data= []  
for i in range(B): 
     new_data.append(tf.gather(data[i], indices[i])) 
new_data= tf.stack(new_data) 

不過,我相信這是不是這樣做的正確方法。有誰知道更好的方法? (我想我應該用tf.gather_nd()以某種方式,但我不知道如何)

我已經看到類似問題的幾個答案here。然而,我找不到解決我的問題。

回答

0

您可以使用tf.gather_nd()有這樣的代碼:

import tensorflow as tf 

# B = 3 
# N = 4 
# M = 2 
# [B x N x 3] 
data = tf.constant([ 
    [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]], 
    [[100, 101, 102], [103, 104, 105], [106, 107, 108], [109, 110, 111]], 
    [[200, 201, 202], [203, 204, 205], [206, 207, 208], [209, 210, 211]], 
    ]) 

# [B x M] 
indices = tf.constant([ 
    [0, 2], 
    [1, 3], 
    [3, 2], 
    ]) 

indices_shape = tf.shape(indices) 

indices_help = tf.tile(tf.reshape(tf.range(indices_shape[0]), [indices_shape[0], 1]) ,[1, indices_shape[1]]); 
indices_ext = tf.concat([tf.expand_dims(indices_help, 2), tf.expand_dims(indices, 2)], axis = 2) 
new_data = tf.gather_nd(data, indices_ext) 

with tf.Session() as sess: 
    sess.run(tf.global_variables_initializer()) 
    print('data') 
    print(sess.run(data)) 
    print('\nindices') 
    print(sess.run(indices)) 
    print('\nnew_data') 
    print(sess.run(new_data)) 

new_data將是:

[[[ 0 1 2] 
    [ 6 7 8]] 

[[103 104 105] 
    [109 110 111]] 

[[209 210 211] 
    [206 207 208]]] 
+0

它的工作原理!但是你能解釋一下這三行代碼的邏輯是什麼嗎('indices_help','indices_ext','new_data')? –

+0

'tf.gather_nd'需要切片索引(兩個座標爲''),但我們只有一個座標'',在索引數組中有隱式座標'b'。因此,我們應該通過添加_batch_的索引來擴展'indices'的每個元素。就是這個。首先我們製作索引爲「批次」的「索引_幫助」,然後將其與源'索引'連接起來 –